fix: Properly clean up state when executing a command fails (#6437)

* fix: Properly clean up state when executing a command fails

ksqldb-rest-app/src/main/java/io/confluent/ksql/rest/server/computation/CommandQueue.java

@@ -87,6 +87,12 @@ void ensureConsumedPast(long seqNum, Duration timeout)
    */
   Producer<CommandId, Command> createTransactionalProducer();
 
+  /**
+   * If a command created with enqueueCommand has had errors and the transaction has failed or been
+   * aborted, this should be called to ensure state is cleaned up.
+   */
+  void abortCommand(CommandId commandId);
+
   /**
    * Blocks until the command topic consumer has processed all records up to
    * the current offset when this method is called.

ksqldb-rest-app/src/main/java/io/confluent/ksql/rest/server/computation/CommandStore.java

@@ -303,6 +303,19 @@ public Producer<CommandId, Command> createTransactionalProducer() {
     );
   }
 
+  @Override
+  public void abortCommand(final CommandId commandId) {
+    commandStatusMap.compute(
+        commandId,
+        (k, v) -> {
+          if (v != null) {
+            LOG.info("Aborting existing command {}", commandId);
+          }
+          return null;
+        }
+    );
+  }
+
   @Override
   public void waitForCommandConsumer() {
     try {

ksqldb-rest-app/src/main/java/io/confluent/ksql/rest/server/computation/DistributingExecutor.java

@@ -135,12 +135,13 @@ public Optional<KsqlEntity> execute(
           "Could not write the statement '%s' into the command topic: " + e.getMessage(),
           statement.getStatementText()), e);
     }
-
+
+    CommandId commandId = null;
     try {
       transactionalProducer.beginTransaction();
       commandQueue.waitForCommandConsumer();
 
-      final CommandId commandId = commandIdAssigner.getCommandId(statement.getStatement());
+      commandId = commandIdAssigner.getCommandId(statement.getStatement());
       final Command command = validatedCommandFactory.create(
           injected,
           executionContext.createSandbox(executionContext.getServiceContext())
@@ -164,11 +165,17 @@ public Optional<KsqlEntity> execute(
     ) {
       // We can't recover from these exceptions, so our only option is close producer and exit.
       // This catch doesn't abortTransaction() since doing that would throw another exception.
+      if (commandId != null) {
+        commandQueue.abortCommand(commandId);
+      }
       throw new KsqlServerException(String.format(
           "Could not write the statement '%s' into the command topic.",
           statement.getStatementText()), e);
     } catch (final Exception e) {
       transactionalProducer.abortTransaction();
+      if (commandId != null) {
+        commandQueue.abortCommand(commandId);
+      }
       throw new KsqlServerException(String.format(
           "Could not write the statement '%s' into the command topic.",
           statement.getStatementText()), e);

ksqldb-rest-app/src/test/java/io/confluent/ksql/rest/server/computation/CommandStoreTest.java

@@ -368,6 +368,16 @@ public void shouldStartCommandTopicOnStart() {
     verify(commandTopic).start();
   }
 
+  @Test
+  public void shouldSuccessfullyAbortAndRetry() {
+    // Given:
+    commandStore.enqueueCommand(commandId, command, transactionalProducer);
+
+    // When/Then:
+    commandStore.abortCommand(commandId);
+    commandStore.enqueueCommand(commandId, command, transactionalProducer);
+  }
+
   private static ConsumerRecords<byte[], byte[]> buildRecords(final Object... args) {
     assertThat(args.length % 2, equalTo(0));
     final List<ConsumerRecord<byte[], byte[]>> records = new ArrayList<>();

ksqldb-rest-app/src/test/java/io/confluent/ksql/rest/server/computation/DistributingExecutorTest.java

@@ -73,6 +73,8 @@
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.function.Supplier;
 import org.apache.kafka.clients.producer.Producer;
+import org.apache.kafka.common.errors.AuthorizationException;
+import org.apache.kafka.common.errors.ProducerFencedException;
 import org.apache.kafka.common.errors.TimeoutException;
 import org.junit.Before;
 import org.junit.Test;
@@ -401,4 +403,36 @@ public void shouldThrowExceptionWhenInsertIntoProcessingLogTopic() {
         "Cannot insert into read-only topic: "
             + "default_ksql_processing_log"));
   }
+
+  @Test
+  public void shouldAbortOnError_ProducerFencedException() {
+    // When:
+    doThrow(new ProducerFencedException("Error!")).when(transactionalProducer).commitTransaction();
+    final Exception e = assertThrows(
+        KsqlServerException.class,
+        () -> distributor.execute(CONFIGURED_STATEMENT, executionContext, securityContext)
+    );
+
+    assertThat(e.getMessage(), containsString("Could not write the statement "
+        + "'statement' into the command topic."));
+
+    // Then:
+    verify(queue).abortCommand(IDGEN.getCommandId(CONFIGURED_STATEMENT.getStatement()));
+  }
+
+  @Test
+  public void shouldAbortOnError_Exception() {
+    // When:
+    doThrow(new RuntimeException("Error!")).when(transactionalProducer).commitTransaction();
+    final Exception e = assertThrows(
+        KsqlServerException.class,
+        () -> distributor.execute(CONFIGURED_STATEMENT, executionContext, securityContext)
+    );
+
+    assertThat(e.getMessage(), containsString("Could not write the statement "
+        + "'statement' into the command topic."));
+
+    // Then:
+    verify(queue).abortCommand(IDGEN.getCommandId(CONFIGURED_STATEMENT.getStatement()));
+  }
 }

ksqldb-rest-app/src/test/java/io/confluent/ksql/rest/server/computation/RecoveryTest.java

@@ -173,6 +173,10 @@ public Producer<CommandId, Command> createTransactionalProducer() {
       return transactionalProducer;
     }
 
+    @Override
+    public void abortCommand(CommandId commandId) {
+    }
+
     @Override
     public boolean corruptionDetected() {
       return false;