fix kafka consumer concurrent access during shutdown (#3193)

extensions-core/kafka-indexing-service/src/main/java/io/druid/indexing/kafka/supervisor/KafkaSupervisor.java

@@ -91,6 +91,7 @@ public class KafkaSupervisor implements Supervisor
   private static final EmittingLogger log = new EmittingLogger(KafkaSupervisor.class);
   private static final Random RANDOM = new Random();
   private static final long MAX_RUN_FREQUENCY_MILLIS = 1000; // prevent us from running too often in response to events
+  private static final int SHUTDOWN_TIMEOUT_MILLIS = 30000;
 
   // Internal data structures
   // --------------------------------------------------------
@@ -166,6 +167,7 @@ public class KafkaSupervisor implements Supervisor
   private final ScheduledExecutorService scheduledExec;
   private final BlockingQueue<Notice> notices = new LinkedBlockingDeque<>();
   private final Object stopLock = new Object();
+  private final Object stateChangeLock = new Object();
 
   private boolean listenerRegistered = false;
   private long lastRunTime;
@@ -238,6 +240,7 @@ public class KafkaSupervisor implements Supervisor
   @Override
   public void start()
   {
+    synchronized (stateChangeLock) {
       Preconditions.checkState(!started, "already started");
       Preconditions.checkState(!exec.isShutdown(), "already stopped");
 
@@ -288,10 +291,12 @@ public class KafkaSupervisor implements Supervisor
       started = true;
       log.info("Started KafkaSupervisor[%s], first run in [%s]", dataSource, ioConfig.getStartDelay());
     }
+  }
 
   @Override
   public void stop(boolean stopGracefully)
   {
+    synchronized (stateChangeLock) {
       Preconditions.checkState(started, "not started");
 
       log.info("Beginning shutdown of KafkaSupervisor[%s]", dataSource);
@@ -308,19 +313,29 @@ public class KafkaSupervisor implements Supervisor
         // until the tasks have acknowledged or timed out. We want this behavior when we're explicitly shut down through
         // the API, but if we shut down for other reasons (e.g. we lose leadership) we want to just stop and leave the
         // tasks as they are.
-      if (stopGracefully) {
-        log.info("Stopping gracefully, signalling managed tasks to complete and publish");
         synchronized (stopLock) {
+          if (stopGracefully) {
+            log.info("Posting GracefulShutdownNotice, signalling managed tasks to complete and publish");
+            notices.add(new GracefulShutdownNotice());
+          } else {
+            log.info("Posting ShutdownNotice");
             notices.add(new ShutdownNotice());
+          }
+
+          long endTime = System.currentTimeMillis() + SHUTDOWN_TIMEOUT_MILLIS;
           while (!stopped) {
-            stopLock.wait();
+            long sleepTime = endTime - System.currentTimeMillis();
+            if (sleepTime <= 0) {
+              log.info("Timed out while waiting for shutdown");
+              stopped = true;
+              break;
+            }
+            stopLock.wait(sleepTime);
           }
         }
         log.info("Shutdown notice handled");
-      }
 
         exec.shutdownNow();
-      consumer.close();
         started = false;
 
         log.info("KafkaSupervisor[%s] has stopped", dataSource);
@@ -330,6 +345,7 @@ public class KafkaSupervisor implements Supervisor
            .emit();
       }
     }
+  }
 
   @Override
   public SupervisorReport getStatus()
@@ -394,7 +410,7 @@ public class KafkaSupervisor implements Supervisor
     }
   }
 
-  private class ShutdownNotice implements Notice
+  private class GracefulShutdownNotice extends ShutdownNotice
   {
     @Override
     public void handle()
@@ -414,6 +430,16 @@ public class KafkaSupervisor implements Supervisor
       }
 
       checkTaskDuration();
+      super.handle();
+    }
+  }
+
+  private class ShutdownNotice implements Notice
+  {
+    @Override
+    public void handle()
+    {
+      consumer.close();
 
       synchronized (stopLock) {
         stopped = true;