Fix potential null pointer issues in StreamOperatorTask

Author: bharathkk <[email protected]> Reviewers: Prateek Maheshwari <[email protected]>, Shanthoosh Venkatraman <[email protected]> Closes apache#699 from bharathkk/bug-fix
shanthoosh · Oct 8, 2018 · e312bb5 · e312bb5
1 parent 2f1003b
commit e312bb5
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 1 deletion.
diff --git a/samza-core/src/main/java/org/apache/samza/task/StreamOperatorTask.java b/samza-core/src/main/java/org/apache/samza/task/StreamOperatorTask.java
@@ -136,7 +136,9 @@ public void close() throws Exception {
     if (this.contextManager != null) {
       this.contextManager.close();
     }
-    operatorImplGraph.close();
+    if (operatorImplGraph != null) {
+      operatorImplGraph.close();
+    }
   }
 
   /* package private for testing */

diff --git a/samza-core/src/test/java/org/apache/samza/task/TestStreamOperatorTask.java b/samza-core/src/test/java/org/apache/samza/task/TestStreamOperatorTask.java
@@ -19,12 +19,37 @@
 
 package org.apache.samza.task;
 
+import org.apache.samza.config.Config;
+import org.apache.samza.operators.ContextManager;
+import org.apache.samza.operators.OperatorSpecGraph;
 import org.apache.samza.operators.impl.OperatorImplGraph;
+import org.junit.Test;
+
+import static org.junit.Assert.*;
+import static org.mockito.Mockito.*;
 
 
 public class TestStreamOperatorTask {
 
   public static OperatorImplGraph getOperatorImplGraph(StreamOperatorTask task) {
     return task.getOperatorImplGraph();
   }
+
+  @Test
+  public void testCloseDuringInitializationErrors() {
+    ContextManager mockContextManager = mock(ContextManager.class);
+    StreamOperatorTask operatorTask = new StreamOperatorTask(mock(OperatorSpecGraph.class), mockContextManager);
+
+    doThrow(new RuntimeException("Failed to initialize context manager"))
+        .when(mockContextManager).init(any(), any());
+
+    try {
+      operatorTask.init(mock(Config.class), mock(TaskContext.class));
+      operatorTask.close();
+    } catch (Exception e) {
+      if (e instanceof NullPointerException) {
+        fail("Unexpected null pointer exception");
+      }
+    }
+  }
 }