Sample: Auto Heartbeat via Interceptor

Signed-off-by: Tihomir Surdilovic <tihomir@temporal.io>

Author: tsurdilo

core/src/main/java/io/temporal/samples/autoheartbeat/AutoHeartbeatUtil.java

@@ -0,0 +1,87 @@
+/*
+ *  Copyright (c) 2020 Temporal Technologies, Inc. All Rights Reserved
+ *
+ *  Copyright 2012-2016 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ *  Modifications copyright (C) 2017 Uber Technologies, Inc.
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License"). You may not
+ *  use this file except in compliance with the License. A copy of the License is
+ *  located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ *  or in the "license" file accompanying this file. This file is distributed on
+ *  an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ *  express or implied. See the License for the specific language governing
+ *  permissions and limitations under the License.
+ */
+
+package io.temporal.samples.autoheartbeat;
+
+import io.temporal.activity.ActivityExecutionContext;
+import java.time.Instant;
+import java.time.ZoneId;
+import java.time.format.DateTimeFormatter;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ScheduledFuture;
+import java.util.concurrent.TimeUnit;
+
+public class AutoHeartbeatUtil {
+  private final long period;
+  private final long initialDelay;
+  private final TimeUnit periodTimeUnit;
+  private final ScheduledExecutorService timerService =
+      Executors.newSingleThreadScheduledExecutor();
+  private final ActivityExecutionContext context;
+  private final Object details;
+  private String heartbeaterId;
+
+  public AutoHeartbeatUtil(
+      long period,
+      long initialDelay,
+      TimeUnit periodTimeUnit,
+      ActivityExecutionContext context,
+      Object details) {
+    this.period = period;
+    this.initialDelay = initialDelay;
+    this.periodTimeUnit = periodTimeUnit;
+    this.context = context;
+    this.details = details;
+    // Set to activity id better, for sample we just use type
+    heartbeaterId = context.getInfo().getActivityType();
+  }
+
+  public ScheduledFuture<?> start() {
+    System.out.println("Autoheartbeater[" + heartbeaterId + "] starting...");
+    return timerService.scheduleAtFixedRate(
+        () -> {
+          //          try {
+          System.out.println(
+              "Autoheartbeater["
+                  + heartbeaterId
+                  + "]"
+                  + "heartbeating at: "
+                  + printShortCurrentTime());
+          context.heartbeat(details);
+        },
+        initialDelay,
+        period,
+        periodTimeUnit);
+  }
+
+  public void stop() {
+    System.out.println("Autoheartbeater[" + heartbeaterId + "] being requested to stop.");
+    // Try not to execute another heartbeat that could have been queued up
+    // Note this can at times take a second or two so make sure to test this out on your workers
+    // So can set best heartbeat timeout (sometimes might need larger value to accomodate)
+    timerService.shutdownNow();
+  }
+
+  private String printShortCurrentTime() {
+    return DateTimeFormatter.ofPattern("HH:mm:ss")
+        .withZone(ZoneId.systemDefault())
+        .format(Instant.now());
+  }
+}

core/src/main/java/io/temporal/samples/autoheartbeat/README.md

@@ -0,0 +1,16 @@
+# Auto-heartbeating sample for activities that define HeartbeatTimeout
+
+This sample shows an implementation of an "auto-heartbeating" utility that can be applied via interceptor to all
+activities where you define HeartbeatTimeout. Use case where this can be helpful include situations where you have
+long-running activities where you want to heartbeat but its difficult to explicitly call heartbeat api in activity code
+directly.
+Another useful scenario for this is where you have activity that at times can complete in very short amount of time,
+but then at times can take for example minutes. In this case you have to set longer StartToClose timeout 
+but you might not want first heartbeat to be sent right away but send it after the "shorter" duration of activity
+execution. 
+
+1. Start the Sample:
+
+```bash
+./gradlew -q execute -PmainClass=io.temporal.samples.autoheartbeat.Starter
+```

core/src/main/java/io/temporal/samples/autoheartbeat/Starter.java

@@ -0,0 +1,103 @@
+/*
+ *  Copyright (c) 2020 Temporal Technologies, Inc. All Rights Reserved
+ *
+ *  Copyright 2012-2016 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ *  Modifications copyright (C) 2017 Uber Technologies, Inc.
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License"). You may not
+ *  use this file except in compliance with the License. A copy of the License is
+ *  located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ *  or in the "license" file accompanying this file. This file is distributed on
+ *  an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ *  express or implied. See the License for the specific language governing
+ *  permissions and limitations under the License.
+ */
+
+package io.temporal.samples.autoheartbeat;
+
+import io.temporal.client.WorkflowClient;
+import io.temporal.client.WorkflowOptions;
+import io.temporal.client.WorkflowStub;
+import io.temporal.samples.autoheartbeat.activities.AutoActivitiesImpl;
+import io.temporal.samples.autoheartbeat.interceptor.AutoHeartbeatWorkerInterceptor;
+import io.temporal.samples.autoheartbeat.workflows.AutoWorkflow;
+import io.temporal.samples.autoheartbeat.workflows.AutoWorkflowImpl;
+import io.temporal.serviceclient.WorkflowServiceStubs;
+import io.temporal.worker.Worker;
+import io.temporal.worker.WorkerFactory;
+import io.temporal.worker.WorkerFactoryOptions;
+
+public class Starter {
+  static final String TASK_QUEUE = "AutoheartbeatTaskQueue";
+  static final String WORKFLOW_ID = "AutoHeartbeatWorkflow";
+
+  public static void main(String[] args) {
+    WorkflowServiceStubs service = WorkflowServiceStubs.newLocalServiceStubs();
+    WorkflowClient client = WorkflowClient.newInstance(service);
+
+    // Configure our auto heartbeat workflow interceptor which will apply
+    // AutoHeartbeaterUtil to each activity workflow schedules which has a heartbeat
+    // timeout configured
+    WorkerFactoryOptions wfo =
+        WorkerFactoryOptions.newBuilder()
+            .setWorkerInterceptors(new AutoHeartbeatWorkerInterceptor())
+            .build();
+
+    WorkerFactory factory = WorkerFactory.newInstance(client, wfo);
+    Worker worker = factory.newWorker(TASK_QUEUE);
+
+    worker.registerWorkflowImplementationTypes(AutoWorkflowImpl.class);
+    worker.registerActivitiesImplementations(new AutoActivitiesImpl());
+
+    factory.start();
+
+    System.out.println("**** First Run: run workflow to completion");
+
+    AutoWorkflow firstRun =
+        client.newWorkflowStub(
+            AutoWorkflow.class,
+            WorkflowOptions.newBuilder()
+                .setWorkflowId(WORKFLOW_ID)
+                .setTaskQueue(TASK_QUEUE)
+                .build());
+
+    try {
+      String firstRunResult = firstRun.exec("Auto heartbeating is cool");
+      System.out.println("First run result: " + firstRunResult);
+    } catch (Exception e) {
+      System.out.println("First run - Workflow exec exception: " + e.getClass().getName());
+    }
+
+    System.out.println("\n\n**** Second Run: cancel activities");
+
+    AutoWorkflow secondRun =
+        client.newWorkflowStub(
+            AutoWorkflow.class,
+            WorkflowOptions.newBuilder()
+                .setWorkflowId(WORKFLOW_ID)
+                .setTaskQueue(TASK_QUEUE)
+                .build());
+    WorkflowClient.start(secondRun::exec, "Auto heartbeating is cool");
+    doSleeps(4);
+    secondRun.cancelActivity();
+
+    try {
+      String secondRunResult = WorkflowStub.fromTyped(secondRun).getResult(String.class);
+      System.out.println("Second run result: " + secondRunResult);
+    } catch (Exception e) {
+      System.out.println("Second run - Workflow exec exception: " + e.getClass().getName());
+    }
+  }
+
+  private static void doSleeps(int seconds) {
+    try {
+      Thread.sleep(seconds * 1000L);
+    } catch (Exception e) {
+      System.out.println(e.getMessage());
+    }
+  }
+}

core/src/main/java/io/temporal/samples/autoheartbeat/activities/AutoActivities.java

@@ -0,0 +1,29 @@
+/*
+ *  Copyright (c) 2020 Temporal Technologies, Inc. All Rights Reserved
+ *
+ *  Copyright 2012-2016 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ *  Modifications copyright (C) 2017 Uber Technologies, Inc.
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License"). You may not
+ *  use this file except in compliance with the License. A copy of the License is
+ *  located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ *  or in the "license" file accompanying this file. This file is distributed on
+ *  an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ *  express or implied. See the License for the specific language governing
+ *  permissions and limitations under the License.
+ */
+
+package io.temporal.samples.autoheartbeat.activities;
+
+import io.temporal.activity.ActivityInterface;
+
+@ActivityInterface
+public interface AutoActivities {
+  String runActivityOne(String input);
+
+  String runActivityTwo(String input);
+}

core/src/main/java/io/temporal/samples/autoheartbeat/activities/AutoActivitiesImpl.java

@@ -0,0 +1,51 @@
+/*
+ *  Copyright (c) 2020 Temporal Technologies, Inc. All Rights Reserved
+ *
+ *  Copyright 2012-2016 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ *  Modifications copyright (C) 2017 Uber Technologies, Inc.
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License"). You may not
+ *  use this file except in compliance with the License. A copy of the License is
+ *  located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ *  or in the "license" file accompanying this file. This file is distributed on
+ *  an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ *  express or implied. See the License for the specific language governing
+ *  permissions and limitations under the License.
+ */
+
+package io.temporal.samples.autoheartbeat.activities;
+
+import java.util.concurrent.TimeUnit;
+
+public class AutoActivitiesImpl implements AutoActivities {
+
+  @Override
+  public String runActivityOne(String input) {
+    return runActivity("runActivityOne - " + input, 10);
+  }
+
+  @Override
+  public String runActivityTwo(String input) {
+    return runActivity("runActivityTwo - " + input, 5);
+  }
+
+  @SuppressWarnings("FutureReturnValueIgnored")
+  private String runActivity(String input, int seconds) {
+    for (int i = 0; i < seconds; i++) {
+      sleep(1);
+    }
+    return "Activity completed: " + input;
+  }
+
+  private void sleep(int seconds) {
+    try {
+      Thread.sleep(TimeUnit.SECONDS.toMillis(seconds));
+    } catch (InterruptedException ee) {
+      // Empty
+    }
+  }
+}

core/src/main/java/io/temporal/samples/autoheartbeat/interceptor/AutoHeartbeatActivityInboundCallsInterceptor.java

@@ -0,0 +1,106 @@
+/*
+ *  Copyright (c) 2020 Temporal Technologies, Inc. All Rights Reserved
+ *
+ *  Copyright 2012-2016 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ *  Modifications copyright (C) 2017 Uber Technologies, Inc.
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License"). You may not
+ *  use this file except in compliance with the License. A copy of the License is
+ *  located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ *  or in the "license" file accompanying this file. This file is distributed on
+ *  an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ *  express or implied. See the License for the specific language governing
+ *  permissions and limitations under the License.
+ */
+
+package io.temporal.samples.autoheartbeat.interceptor;
+
+import io.temporal.activity.Activity;
+import io.temporal.activity.ActivityExecutionContext;
+import io.temporal.client.ActivityCanceledException;
+import io.temporal.common.interceptors.ActivityInboundCallsInterceptor;
+import io.temporal.common.interceptors.ActivityInboundCallsInterceptorBase;
+import io.temporal.samples.autoheartbeat.AutoHeartbeatUtil;
+import java.time.Duration;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ScheduledFuture;
+import java.util.concurrent.TimeUnit;
+
+public class AutoHeartbeatActivityInboundCallsInterceptor
+    extends ActivityInboundCallsInterceptorBase {
+  private ActivityExecutionContext activityExecutionContext;
+  private Duration activityHeartbeatTimeout;
+  private AutoHeartbeatUtil autoHeartbeater;
+  //  private CompletableFuture autoHeartbeatFuture;
+  private ScheduledFuture scheduledFuture;
+
+  //  private ScheduledFuture scheduledFuture;
+
+  public AutoHeartbeatActivityInboundCallsInterceptor(ActivityInboundCallsInterceptor next) {
+    super(next);
+  }
+
+  @Override
+  public void init(ActivityExecutionContext context) {
+    this.activityExecutionContext = context;
+    activityHeartbeatTimeout = activityExecutionContext.getInfo().getHeartbeatTimeout();
+    super.init(context);
+  }
+
+  @Override
+  @SuppressWarnings({"FutureReturnValueIgnored", "CatchAndPrintStackTrace"})
+  public ActivityOutput execute(ActivityInput input) {
+    // If activity has heartbeat timeout defined we want to apply auto-heartbeter
+    if (activityHeartbeatTimeout != null && activityHeartbeatTimeout.getSeconds() > 0) {
+      System.out.println(
+          "Auto heartbeating applied for activity: "
+              + activityExecutionContext.getInfo().getActivityType());
+      autoHeartbeater =
+          new AutoHeartbeatUtil(2, 0, TimeUnit.SECONDS, activityExecutionContext, input);
+      scheduledFuture = autoHeartbeater.start();
+    } else {
+      System.out.println(
+          "Auto heartbeating not being applied for activity: "
+              + activityExecutionContext.getInfo().getActivityType());
+    }
+
+    if (scheduledFuture != null) {
+      CompletableFuture activityExecFuture =
+          CompletableFuture.supplyAsync(() -> super.execute(input));
+      CompletableFuture autoHeartbeatFuture =
+          CompletableFuture.supplyAsync(
+              () -> {
+                try {
+                  return scheduledFuture.get();
+                } catch (Exception e) {
+                  throw new ActivityCanceledException(activityExecutionContext.getInfo());
+                }
+              });
+      try {
+        CompletableFuture.anyOf(autoHeartbeatFuture, activityExecFuture).get();
+      } catch (Exception e) {
+        if (e instanceof ExecutionException) {
+          ExecutionException ee = (ExecutionException) e;
+          if (ee.getCause() instanceof ActivityCanceledException) {
+            throw new ActivityCanceledException(activityExecutionContext.getInfo());
+          }
+        }
+        throw Activity.wrap(e);
+      } finally {
+        if (autoHeartbeater != null) {
+          autoHeartbeater.stop();
+        }
+      }
+    }
+    return super.execute(input);
+  }
+
+  public interface AutoHeartbeaterCancellationCallback {
+    void handle(Exception e);
+  }
+}