Cancel the existing grok task when a timeout occurs. Resolves #4026 (#…

…4027) (#4028) Signed-off-by: David Venable <[email protected]> (cherry picked from commit 7616396) Co-authored-by: David Venable <[email protected]>
opensearch-project · Jan 30, 2024 · 5d25021 · 5d25021
1 parent f7b9b4d
commit 5d25021
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 5 deletions.
diff --git a/...cessor/src/main/java/org/opensearch/dataprepper/plugins/processor/grok/GrokProcessor.java b/...cessor/src/main/java/org/opensearch/dataprepper/plugins/processor/grok/GrokProcessor.java
@@ -319,7 +319,12 @@ private boolean shouldBreakOnMatch(final Map<String, Object> captures) {
     }
 
     private void runWithTimeout(final Runnable runnable) throws TimeoutException, ExecutionException, InterruptedException {
-        Future<?> task = executorService.submit(runnable);
-        task.get(grokProcessorConfig.getTimeoutMillis(), TimeUnit.MILLISECONDS);
+        final Future<?> task = executorService.submit(runnable);
+        try {
+            task.get(grokProcessorConfig.getTimeoutMillis(), TimeUnit.MILLISECONDS);
+        } catch (final TimeoutException exception) {
+            task.cancel(true);
+            throw exception;
+        }
     }
 }
diff --git a/...r/src/test/java/org/opensearch/dataprepper/plugins/processor/grok/GrokProcessorTests.java b/...r/src/test/java/org/opensearch/dataprepper/plugins/processor/grok/GrokProcessorTests.java
@@ -52,6 +52,7 @@
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.verifyNoInteractions;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
 import static org.mockito.Mockito.when;
 import static org.opensearch.dataprepper.plugins.processor.grok.GrokProcessor.EXECUTOR_SERVICE_SHUTDOWN_TIMEOUT;
 import static org.opensearch.dataprepper.test.matcher.MapEquals.isEqualWithoutTimestamp;
@@ -155,7 +156,7 @@ private GrokProcessor createObjectUnderTest() {
     }
 
     @Test
-    public void testMatchMerge() throws JsonProcessingException {
+    public void testMatchMerge() throws JsonProcessingException, ExecutionException, InterruptedException, TimeoutException {
         grokProcessor = createObjectUnderTest();
 
         capture.put("key_capture_1", "value_capture_1");
@@ -182,10 +183,12 @@ public void testMatchMerge() throws JsonProcessingException {
         verify(grokProcessingMatchCounter, times(1)).increment();
         verify(grokProcessingTime, times(1)).record(any(Runnable.class));
         verifyNoInteractions(grokProcessingErrorsCounter, grokProcessingMismatchCounter, grokProcessingTimeoutsCounter);
+        verify(task).get(GrokProcessorConfig.DEFAULT_TIMEOUT_MILLIS, TimeUnit.MILLISECONDS);
+        verifyNoMoreInteractions(task);
     }
 
     @Test
-    public void testTarget() throws JsonProcessingException {
+    public void testTarget() throws JsonProcessingException, ExecutionException, InterruptedException, TimeoutException {
         pluginSetting.getSettings().put(GrokProcessorConfig.TARGET_KEY, "test_target");
         grokProcessor = createObjectUnderTest();
 
@@ -216,6 +219,8 @@ public void testTarget() throws JsonProcessingException {
         verify(grokProcessingMatchCounter, times(1)).increment();
         verify(grokProcessingTime, times(1)).record(any(Runnable.class));
         verifyNoInteractions(grokProcessingErrorsCounter, grokProcessingMismatchCounter, grokProcessingTimeoutsCounter);
+        verify(task).get(GrokProcessorConfig.DEFAULT_TIMEOUT_MILLIS, TimeUnit.MILLISECONDS);
+        verifyNoMoreInteractions(task);
     }
 
     @Test
@@ -251,7 +256,7 @@ public void testOverwrite() throws JsonProcessingException {
     }
 
     @Test
-    public void testMatchMergeCollisionStrings() throws JsonProcessingException {
+    public void testMatchMergeCollisionStrings() throws JsonProcessingException, ExecutionException, InterruptedException, TimeoutException {
         grokProcessor = createObjectUnderTest();
 
         capture.put("key_capture_1", "value_capture_1");
@@ -398,6 +403,7 @@ public void testThatTimeoutExceptionIsCaughtAndProcessingContinues() throws Json
         assertThat(grokkedRecords.size(), equalTo(1));
         assertThat(grokkedRecords.get(0), notNullValue());
         assertRecordsAreEqual(grokkedRecords.get(0), record);
+        verify(task).cancel(true);
         verify(grokProcessingTimeoutsCounter, times(1)).increment();
         verify(grokProcessingTime, times(1)).record(any(Runnable.class));
     }