Merge branch 'master' of git://github.com/lyft/amazon-kinesis-client into lyft-master

src/main/java/com/amazonaws/services/kinesis/multilang/MessageWriter.java

@@ -26,13 +26,16 @@ import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 
 import com.amazonaws.services.kinesis.clientlibrary.lib.worker.ShutdownReason;
-import com.amazonaws.services.kinesis.clientlibrary.types.InitializationInput;
+
-import com.amazonaws.services.kinesis.clientlibrary.types.ProcessRecordsInput;
 import com.amazonaws.services.kinesis.multilang.messages.CheckpointMessage;
 import com.amazonaws.services.kinesis.multilang.messages.InitializeMessage;
 import com.amazonaws.services.kinesis.multilang.messages.Message;
 import com.amazonaws.services.kinesis.multilang.messages.ProcessRecordsMessage;
 import com.amazonaws.services.kinesis.multilang.messages.ShutdownMessage;
+import com.amazonaws.services.kinesis.multilang.messages.ShutdownRequestedMessage;
+
+import com.amazonaws.services.kinesis.clientlibrary.types.InitializationInput;
+import com.amazonaws.services.kinesis.clientlibrary.types.ProcessRecordsInput;
 import com.fasterxml.jackson.databind.ObjectMapper;
 
 /**
@@ -145,6 +148,13 @@ class MessageWriter {
         return writeMessage(new ShutdownMessage(reason));
     }
 
+    /**
+     * Writes a {@link ShutdownRequestedMessage} to the subprocess.
+     */
+    Future<Boolean> writeShutdownRequestedMessage() {
+        return writeMessage(new ShutdownRequestedMessage());
+    }
+
     /**
      * Writes a {@link CheckpointMessage} to the subprocess.
      * 

src/main/java/com/amazonaws/services/kinesis/multilang/MultiLangDaemon.java

@@ -16,10 +16,13 @@ package com.amazonaws.services.kinesis.multilang;
 
 import java.io.IOException;
 import java.io.PrintStream;
+
 import java.util.concurrent.Callable;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
@@ -140,11 +143,25 @@ public class MultiLangDaemon implements Callable<Integer> {
         ExecutorService executorService = config.getExecutorService();
 
         // Daemon
-        MultiLangDaemon daemon = new MultiLangDaemon(
+        final MultiLangDaemon daemon = new MultiLangDaemon(
                 config.getKinesisClientLibConfiguration(),
                 config.getRecordProcessorFactory(),
                 executorService);
 
+        Runtime.getRuntime().addShutdownHook(new Thread() {
+            @Override
+            public void run() {
+                LOG.info("Process terminanted, will initiate shutdown.");
+                try {
+                    Future<Void> fut = daemon.worker.requestShutdown();
+                    fut.get(5000, TimeUnit.MILLISECONDS);
+                    LOG.info("Process shutdown is complete.");
+                } catch (InterruptedException | ExecutionException | TimeoutException e) {
+                    LOG.error("Encountered an error during shutdown.", e);
+                }
+            }
+        });
+
         Future<Integer> future = executorService.submit(daemon);
         try {
             System.exit(future.get());

src/main/java/com/amazonaws/services/kinesis/multilang/MultiLangProtocol.java

@@ -27,6 +27,7 @@ import com.amazonaws.services.kinesis.multilang.messages.InitializeMessage;
 import com.amazonaws.services.kinesis.multilang.messages.Message;
 import com.amazonaws.services.kinesis.multilang.messages.ProcessRecordsMessage;
 import com.amazonaws.services.kinesis.multilang.messages.ShutdownMessage;
+import com.amazonaws.services.kinesis.multilang.messages.ShutdownRequestedMessage;
 import com.amazonaws.services.kinesis.multilang.messages.StatusMessage;
 
 import lombok.extern.apachecommons.CommonsLog;
@@ -99,6 +100,18 @@ class MultiLangProtocol {
         return waitForStatusMessage(ShutdownMessage.ACTION, checkpointer, writeFuture);
     }
 
+    /**
+     * Writes a {@link ShutdownRequestedMessage} to the child process's STDIN and waits for the child process to respond with a
+     * {@link StatusMessage} on its STDOUT.
+     *
+     * @param checkpointer A checkpointer.
+     * @return Whether or not this operation succeeded.
+     */
+    boolean shutdownRequested(IRecordProcessorCheckpointer checkpointer) {
+        Future<Boolean> writeFuture = messageWriter.writeShutdownRequestedMessage();
+        return waitForStatusMessage(ShutdownRequestedMessage.ACTION, checkpointer, writeFuture);
+    }
+
     /**
      * Waits for a {@link StatusMessage} for a particular action. If a {@link CheckpointMessage} is received, then this
      * method will attempt to checkpoint with the provided {@link IRecordProcessorCheckpointer}. This method returns

src/main/java/com/amazonaws/services/kinesis/multilang/MultiLangRecordProcessor.java

@@ -20,6 +20,10 @@ import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Future;
 
+import com.amazonaws.services.kinesis.clientlibrary.exceptions.InvalidStateException;
+import com.amazonaws.services.kinesis.clientlibrary.exceptions.ShutdownException;
+import com.amazonaws.services.kinesis.clientlibrary.interfaces.IRecordProcessorCheckpointer;
+import com.amazonaws.services.kinesis.clientlibrary.interfaces.v2.IShutdownNotificationAware;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 
@@ -29,13 +33,14 @@ import com.amazonaws.services.kinesis.clientlibrary.types.ProcessRecordsInput;
 import com.amazonaws.services.kinesis.clientlibrary.types.ShutdownInput;
 import com.fasterxml.jackson.databind.ObjectMapper;
 
+
 /**
  * A record processor that manages creating a child process that implements the multi language protocol and connecting
  * that child process's input and outputs to a {@link MultiLangProtocol} object and calling the appropriate methods on
  * that object when its corresponding {@link #initialize}, {@link #processRecords}, and {@link #shutdown} methods are
  * called.
  */
-public class MultiLangRecordProcessor implements IRecordProcessor {
+public class MultiLangRecordProcessor implements IRecordProcessor, IShutdownNotificationAware {
 
     private static final Log LOG = LogFactory.getLog(MultiLangRecordProcessor.class);
     private static final int EXIT_VALUE = 1;
@@ -136,6 +141,20 @@ public class MultiLangRecordProcessor implements IRecordProcessor {
         }
     }
 
+    @Override
+    public void shutdownRequested(IRecordProcessorCheckpointer checkpointer) {
+        LOG.info("Shutdown is requested.");
+        if (!initialized) {
+            LOG.info("Record processor was not initialized so no need to initiate a final checkpoint.");
+            return;
+        }
+        LOG.info("Requesting a checkpoint on shutdown notification.");
+        if (!protocol.shutdownRequested(checkpointer)) {
+            LOG.error("Child process failed to complete shutdown notification.");
+        }
+    }
+
+
     /**
      * Used to tell whether the processor has been shutdown already.
      */

src/main/java/com/amazonaws/services/kinesis/multilang/messages/ShutdownRequestedMessage.java

@@ -0,0 +1,17 @@
+package com.amazonaws.services.kinesis.multilang.messages;
+
+/**
+ * A message to indicate to the client's process that shutdown is requested.
+ */
+public class ShutdownRequestedMessage extends Message {
+    /**
+     * The name used for the action field in {@link Message}.
+     */
+    public static final String ACTION = "shutdownRequested";
+
+    /**
+     * Convenience constructor.
+     */
+    public ShutdownRequestedMessage() {
+    }
+}

src/test/java/com/amazonaws/services/kinesis/multilang/MessageWriterTest.java

@@ -113,6 +113,16 @@ public class MessageWriterTest {
         Mockito.verify(this.stream, Mockito.atLeastOnce()).flush();
     }
 
+    @Test
+    public void writeShutdownRequestedMessageTest() throws IOException, InterruptedException, ExecutionException {
+        Future<Boolean> future = this.messageWriter.writeShutdownRequestedMessage();
+        future.get();
+
+        Mockito.verify(this.stream, Mockito.atLeastOnce()).write(Mockito.any(byte[].class), Mockito.anyInt(),
+                Mockito.anyInt());
+        Mockito.verify(this.stream, Mockito.atLeastOnce()).flush();
+    }
+
     @Test
     public void streamIOExceptionTest() throws IOException, InterruptedException, ExecutionException {
         Mockito.doThrow(IOException.class).when(stream).flush();

src/test/java/com/amazonaws/services/kinesis/multilang/MultiLangProtocolTest.java

@@ -114,6 +114,16 @@ public class MultiLangProtocolTest {
         assertThat(protocol.shutdown(null, ShutdownReason.ZOMBIE), equalTo(true));
     }
 
+    @Test
+    public void shutdownRequestedTest() {
+        when(messageWriter.writeShutdownRequestedMessage()).thenReturn(buildFuture(true));
+        when(messageReader.getNextMessageFromSTDOUT()).thenReturn(buildFuture(new StatusMessage("shutdownRequested"), Message.class));
+        Mockito.doReturn(buildFuture(true)).when(messageWriter)
+                .writeShutdownRequestedMessage();
+        Mockito.doReturn(buildFuture(new StatusMessage("shutdownRequested"))).when(messageReader).getNextMessageFromSTDOUT();
+        assertThat(protocol.shutdownRequested(null), equalTo(true));
+    }
+
     private Answer<Future<Message>> buildMessageAnswers(List<Message> messages) {
         return new Answer<Future<Message>>() {
 

src/test/java/com/amazonaws/services/kinesis/multilang/messages/MessageTest.java

@@ -44,7 +44,7 @@ public class MessageTest {
                                 });
                             }
                         })), new ShutdownMessage(ShutdownReason.ZOMBIE), new StatusMessage("processRecords"),
-                        new InitializeMessage(), new ProcessRecordsMessage() };
+                        new InitializeMessage(), new ProcessRecordsMessage(), new ShutdownRequestedMessage() };
 
         for (int i = 0; i < messages.length; i++) {
             Assert.assertTrue("Each message should contain the action field", messages[i].toString().contains("action"));