Need to reinitialize certain components for externally managed runtimes when...

Need to reinitialize certain components for externally managed runtimes when moving functions (#5007)

(cherry picked from commit 27a9f628)

pulsar-functions/runtime/src/main/java/org/apache/pulsar/functions/runtime/KubernetesRuntime.java

@@ -247,9 +247,19 @@ public class KubernetesRuntime implements Runtime {
             throw e;
         }
 
-        if (channel == null && stub == null) {
+        setupGrpcChannelIfNeeded();
+    }
+
+    @Override
+    public void reinitialize() {
+        setupGrpcChannelIfNeeded();
+    }
+
+    private synchronized void setupGrpcChannelIfNeeded() {
+        if (channel == null || stub == null) {
             channel = new ManagedChannel[instanceConfig.getFunctionDetails().getParallelism()];
             stub = new InstanceControlGrpc.InstanceControlFutureStub[instanceConfig.getFunctionDetails().getParallelism()];
+
             String jobName = createJobName(instanceConfig.getFunctionDetails());
             for (int i = 0; i < instanceConfig.getFunctionDetails().getParallelism(); ++i) {
                 String address = getServiceUrl(jobName, jobNamespace, i);
@@ -332,16 +342,16 @@ public class KubernetesRuntime implements Runtime {
     @Override
     public CompletableFuture<InstanceCommunication.MetricsData> getMetrics(int instanceId) {
         CompletableFuture<InstanceCommunication.MetricsData> retval = new CompletableFuture<>();
-        if (instanceId < 0 || instanceId >= stub.length) {
-            if (stub == null) {
-                retval.completeExceptionally(new RuntimeException("Invalid InstanceId"));
-                return retval;
-            }
-        }
         if (stub == null) {
             retval.completeExceptionally(new RuntimeException("Not alive"));
             return retval;
         }
+
+        if (instanceId < 0 || instanceId >= stub.length) {
+            retval.completeExceptionally(new RuntimeException("Invalid InstanceId"));
+            return retval;
+        }
+
         ListenableFuture<InstanceCommunication.MetricsData> response = stub[instanceId].withDeadlineAfter(GRPC_TIMEOUT_SECS, TimeUnit.SECONDS).getMetrics(Empty.newBuilder().build());
         Futures.addCallback(response, new FutureCallback<InstanceCommunication.MetricsData>() {
             @Override

pulsar-functions/runtime/src/main/java/org/apache/pulsar/functions/runtime/Runtime.java

@@ -31,6 +31,10 @@ public interface Runtime {
 
     void start() throws Exception;
 
+    default void reinitialize() {
+
+    }
+
     void join() throws Exception;
 
     void stop() throws Exception;

pulsar-functions/worker/src/main/java/org/apache/pulsar/functions/worker/FunctionRuntimeManager.java

@@ -215,9 +215,6 @@ public class FunctionRuntimeManager implements AutoCloseable{
                     }
                 }
             }
-            // start assignment tailer
-            this.functionAssignmentTailer.start();
-
         } catch (Exception e) {
             log.error("Failed to initialize function runtime manager: ", e.getMessage(), e);
             throw new RuntimeException(e);
@@ -229,7 +226,6 @@ public class FunctionRuntimeManager implements AutoCloseable{
      */
     public void start() {
         log.info("/** Starting Function Runtime Manager **/");
-        log.info("Initialize metrics sink...");
         log.info("Starting function assignment tailer...");
         this.functionAssignmentTailer.start();
     }
@@ -629,7 +625,6 @@ public class FunctionRuntimeManager implements AutoCloseable{
             // changes to the function meta data of the instance
 
             if (runtimeFactory.externallyManaged()) {
-
                 // change in metadata thus need to potentially restart
                 if (!assignment.getInstance().equals(existingAssignment.getInstance())) {
                     //stop function
@@ -657,6 +652,8 @@ public class FunctionRuntimeManager implements AutoCloseable{
                         RuntimeSpawner runtimeSpawner = functionActioner.getRuntimeSpawner(
                                 assignment.getInstance(),
                                 assignment.getInstance().getFunctionMetaData().getPackageLocation().getPackagePath());
+                        // re-initialize if necessary
+                        runtimeSpawner.getRuntime().reinitialize();
                         newFunctionRuntimeInfo.setRuntimeSpawner(runtimeSpawner);
 
                         this.functionRuntimeInfoMap.put(fullyQualifiedInstanceId, newFunctionRuntimeInfo);

pulsar-functions/worker/src/test/java/org/apache/pulsar/functions/worker/FunctionRuntimeManagerTest.java

@@ -648,7 +648,8 @@ public class FunctionRuntimeManagerTest {
         doNothing().when(kubernetesRuntimeFactory).setupClient();
         doReturn(true).when(kubernetesRuntimeFactory).externallyManaged();
 
-        doReturn(mock(KubernetesRuntime.class)).when(kubernetesRuntimeFactory).createContainer(any(), any(), any(), any());
+        KubernetesRuntime kubernetesRuntime = mock(KubernetesRuntime.class);
+        doReturn(kubernetesRuntime).when(kubernetesRuntimeFactory).createContainer(any(), any(), any(), any());
 
         FunctionActioner functionActioner = spy(new FunctionActioner(
                 workerConfig,
@@ -743,7 +744,10 @@ public class FunctionRuntimeManagerTest {
                 instance.getInstanceId());
         Assert.assertTrue(
                 functionRuntimeManager.functionRuntimeInfoMap.get("test-tenant/test-namespace/func-1:0").getRuntimeSpawner().getRuntimeFactory() instanceof KubernetesRuntimeFactory);
+
         Assert.assertTrue(
                 functionRuntimeManager.functionRuntimeInfoMap.get("test-tenant/test-namespace/func-1:0").getRuntimeSpawner().getRuntime() != null);
+
+        verify(kubernetesRuntime, times(1)).reinitialize();
     }
 }