fix: don't cleanup topics on engine close (#4658)

Co-authored-by: Rohan <[email protected]>
Co-authored-by: Andy Coates <[email protected]>

ksql-engine/src/main/java/io/confluent/ksql/engine/KsqlEngine.java

@@ -190,7 +190,7 @@ public ExecuteResult execute(
 
   @Override
   public void close() {
-    allLiveQueries.forEach(QueryMetadata::close);
+    allLiveQueries.forEach(QueryMetadata::stop);
     engineMetrics.close();
     aggregateMetricsCollector.shutdown();
   }

ksql-engine/src/main/java/io/confluent/ksql/util/PersistentQueryMetadata.java

@@ -120,4 +120,9 @@ public String getSchemasDescription() {
   public PersistenceSchemas getPersistenceSchemas() {
     return persistenceSchemas;
   }
+
+  @Override
+  public void stop() {
+    doClose(false);
+  }
 }

ksql-engine/src/main/java/io/confluent/ksql/util/QueryMetadata.java

@@ -30,7 +30,7 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-public class QueryMetadata {
+public abstract class QueryMetadata {
 
   private static final Logger LOG = LoggerFactory.getLogger(QueryMetadata.class);
 
@@ -148,14 +148,34 @@ public boolean hasEverBeenStarted() {
     return everStarted;
   }
 
+
+  /**
+   * Stops the query without cleaning up the external resources
+   * so that it can be resumed when we call {@link #start()}.
+   *
+   */
+  public abstract void stop();
+
+  /**
+   * Closes the {@code QueryMetadata} and cleans up any of
+   * the resources associated with it (e.g. internal topics,
+   * schemas, etc...).
+   *
+   * @see QueryMetadata#stop()
+   */
   public void close() {
-    kafkaStreams.close();
+    doClose(true);
+    closeCallback.accept(this);
+  }
 
-    kafkaStreams.cleanUp();
+  protected void doClose(final boolean cleanUp) {
+    kafkaStreams.close(Duration.ofMillis(closeTimeout));
 
-    queryStateListener.ifPresent(QueryStateListener::close);
+    if (cleanUp) {
+      kafkaStreams.cleanUp();
+    }
 
-    closeCallback.accept(this);
+    queryStateListener.ifPresent(QueryStateListener::close);
   }
 
   public void start() {

ksql-engine/src/test/java/io/confluent/ksql/engine/KsqlEngineTest.java

@@ -34,6 +34,7 @@
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.atLeastOnce;
 import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.verifyNoMoreInteractions;
 
@@ -470,6 +471,63 @@ public void shouldCleanUpInternalTopicsOnClose() {
     verify(topicClient).deleteInternalTopics(query.getQueryApplicationId());
   }
 
+  @Test
+  public void shouldCleanUpInternalTopicsOnEngineCloseForTransientQueries() {
+    // Given:
+    final QueryMetadata query = KsqlEngineTestUtil.executeQuery(
+        serviceContext,
+        ksqlEngine,
+        "select * from test1 EMIT CHANGES;",
+        KSQL_CONFIG, Collections.emptyMap()
+    );
+
+    query.start();
+
+    // When:
+    ksqlEngine.close();
+
+    // Then:
+    verify(topicClient).deleteInternalTopics(query.getQueryApplicationId());
+  }
+
+  @Test
+  public void shouldNotCleanUpInternalTopicsOnEngineCloseForPersistentQueries() {
+    // Given:
+    final List<QueryMetadata> query = KsqlEngineTestUtil.execute(
+        serviceContext,
+        ksqlEngine,
+        "create stream persistent as select * from test1 EMIT CHANGES;",
+        KSQL_CONFIG, Collections.emptyMap()
+    );
+
+    query.get(0).start();
+
+    // When:
+    ksqlEngine.close();
+
+    // Then (there are no transient queries, so no internal topics should be deleted):
+    verify(topicClient, never()).deleteInternalTopics(any());
+  }
+
+  @Test
+  public void shouldCleanUpInternalTopicsOnQueryCloseForPersistentQueries() {
+    // Given:
+    final List<QueryMetadata> query = KsqlEngineTestUtil.execute(
+        serviceContext,
+        ksqlEngine,
+        "create stream persistent as select * from test1 EMIT CHANGES;",
+        KSQL_CONFIG, Collections.emptyMap()
+    );
+
+    query.get(0).start();
+
+    // When:
+    query.get(0).close();
+
+    // Then (there are no transient queries, so no internal topics should be deleted):
+    verify(topicClient).deleteInternalTopics(query.get(0).getQueryApplicationId());
+  }
+
   @Test
   public void shouldNotCleanUpInternalTopicsOnCloseIfQueryNeverStarted() {
     // Given:

ksql-engine/src/test/java/io/confluent/ksql/util/QueryMetadataTest.java

@@ -19,7 +19,9 @@
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.is;
 import static org.mockito.Mockito.inOrder;
+import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
 import static org.mockito.Mockito.when;
 
 import com.google.common.collect.ImmutableSet;
@@ -57,9 +59,11 @@ public class QueryMetadataTest {
   @Mock
   private Consumer<QueryMetadata> closeCallback;
   private QueryMetadata query;
+  private boolean cleanUp;
 
   @Before
   public void setup() {
+    cleanUp = false;
     query = new QueryMetadata(
         "foo",
         kafkaStreams,
@@ -71,8 +75,13 @@ public void setup() {
         topoplogy,
         Collections.emptyMap(),
         Collections.emptyMap(),
-        closeCallback
-    );
+        closeCallback,
+        closeTimeout) {
+      @Override
+      public void stop() {
+        doClose(cleanUp);
+      }
+    };
   }
 
   @Test
@@ -134,6 +143,24 @@ public void shouldCloseKStreamsAppOnCloseThenCloseCallback() {
     inOrder.verify(closeCallback).accept(query);
   }
 
+  @Test
+  public void shouldNotCallCloseCallbackOnStop() {
+    // When:
+    query.stop();
+
+    // Then:
+    verifyNoMoreInteractions(closeCallback);
+  }
+
+  @Test
+  public void shouldCallKafkaStreamsCloseOnStop() {
+    // When:
+    query.stop();
+
+    // Then:
+    verify(kafkaStreams).close(Duration.ofMillis(closeTimeout));
+  }
+
   @Test
   public void shouldCleanUpKStreamsAppAfterCloseOnClose() {
     // When:
@@ -145,6 +172,27 @@ public void shouldCleanUpKStreamsAppAfterCloseOnClose() {
     inOrder.verify(kafkaStreams).cleanUp();
   }
 
+  @Test
+  public void shouldNotCleanUpKStreamsAppOnStop() {
+    // When:
+    query.stop();
+
+    // Then:
+    verify(kafkaStreams, never()).cleanUp();
+  }
+
+  @Test
+  public void shouldCallCleanupOnStopIfCleanup() {
+    // Given:
+    cleanUp = true;
+
+    // When:
+    query.stop();
+
+    // Then:
+    verify(kafkaStreams).cleanUp();
+  }
+
   @Test
   public void shouldReturnSources() {
     assertThat(query.getSourceNames(), is(SOME_SOURCES));