SAMZA-2751: RunLoop and High-Level API changes for Drain (#1616)
diff --git a/gradle.properties b/gradle.properties
index 1a6ee97..cab76e0 100644
--- a/gradle.properties
+++ b/gradle.properties
@@ -16,7 +16,7 @@
# under the License.
group=org.apache.samza
version=1.7.0-SNAPSHOT
-scalaSuffix=2.11
+scalaSuffix=2.12
# after changing this value, run `$ ./gradlew wrapper` and commit the resulting changed files
gradleVersion=5.2.1
diff --git a/samza-api/src/main/java/org/apache/samza/system/DrainMessage.java b/samza-api/src/main/java/org/apache/samza/system/DrainMessage.java
new file mode 100644
index 0000000..4373d51
--- /dev/null
+++ b/samza-api/src/main/java/org/apache/samza/system/DrainMessage.java
@@ -0,0 +1,72 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.samza.system;
+
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+/**
+ * The DrainMessage is a control message that is sent out to next stage
+ * once the task has consumed to the end of a bounded stream.
+ */
+public class DrainMessage extends ControlMessage {
+ /**
+ * Id used to invalidate DrainMessages between runs. Ties to app.run.id from config.
+ */
+ private final String runId;
+
+ public DrainMessage(String runId) {
+ this(null, runId);
+ }
+
+ public DrainMessage(@JsonProperty("task-name") String taskName, @JsonProperty("run-id") String runId) {
+ super(taskName);
+ this.runId = runId;
+ }
+
+ public String getRunId() {
+ return runId;
+ }
+
+ @Override
+ public int hashCode() {
+ final int prime = 31;
+ final int result = prime * super.hashCode() + (this.runId != null ? this.runId.hashCode() : 0);
+ return result;
+ }
+
+ @Override
+ public boolean equals(Object obj) {
+ if (this == obj)
+ return true;
+ if (obj == null)
+ return false;
+ if (getClass() != obj.getClass())
+ return false;
+
+ final DrainMessage other = (DrainMessage) obj;
+ if (!super.equals(other)) {
+ return false;
+ }
+ if (!runId.equals(other.runId)) {
+ return false;
+ }
+ return true;
+ }
+}
diff --git a/samza-api/src/main/java/org/apache/samza/system/IncomingMessageEnvelope.java b/samza-api/src/main/java/org/apache/samza/system/IncomingMessageEnvelope.java
index 2112463..1f4a740 100644
--- a/samza-api/src/main/java/org/apache/samza/system/IncomingMessageEnvelope.java
+++ b/samza-api/src/main/java/org/apache/samza/system/IncomingMessageEnvelope.java
@@ -132,7 +132,6 @@
if (envelopeKeyorOffset == null) {
return new SystemStreamPartition(systemStreamPartition, 0);
}
-
// modulo 31 first to best spread out the hashcode and then modulo elasticityFactor for actual keyBucket
// Note: elasticityFactor <= 16 so modulo 31 is safe to do.
int keyBucket = (Math.abs(envelopeKeyorOffset.hashCode()) % 31) % elasticityFactor;
@@ -162,6 +161,10 @@
return END_OF_STREAM_OFFSET.equals(offset);
}
+ public boolean isDrain() {
+ return message != null && DrainMessage.class.isAssignableFrom(message.getClass());
+ }
+
/**
* This method is deprecated in favor of WatermarkManager.buildEndOfStreamEnvelope(SystemStreamPartition ssp).
*
@@ -172,6 +175,10 @@
return new IncomingMessageEnvelope(ssp, END_OF_STREAM_OFFSET, null, new EndOfStreamMessage(null));
}
+ public static IncomingMessageEnvelope buildDrainMessage(SystemStreamPartition ssp, String runId) {
+ return new IncomingMessageEnvelope(ssp, null, null, new DrainMessage(runId));
+ }
+
public static IncomingMessageEnvelope buildWatermarkEnvelope(SystemStreamPartition ssp, long watermark) {
return new IncomingMessageEnvelope(ssp, null, null, new WatermarkMessage(watermark, null));
}
diff --git a/samza-api/src/main/java/org/apache/samza/system/MessageType.java b/samza-api/src/main/java/org/apache/samza/system/MessageType.java
index 7129d00..9f58621 100644
--- a/samza-api/src/main/java/org/apache/samza/system/MessageType.java
+++ b/samza-api/src/main/java/org/apache/samza/system/MessageType.java
@@ -26,7 +26,8 @@
public enum MessageType {
USER_MESSAGE,
WATERMARK,
- END_OF_STREAM;
+ END_OF_STREAM,
+ DRAIN;
/**
* Returns the {@link MessageType} of a particular intermediate stream message.
@@ -38,6 +39,8 @@
return WATERMARK;
} else if (message instanceof EndOfStreamMessage) {
return END_OF_STREAM;
+ } else if (message instanceof DrainMessage) {
+ return DRAIN;
} else {
return USER_MESSAGE;
}
diff --git a/samza-api/src/main/java/org/apache/samza/task/DrainListenerTask.java b/samza-api/src/main/java/org/apache/samza/task/DrainListenerTask.java
new file mode 100644
index 0000000..c5ca7d5
--- /dev/null
+++ b/samza-api/src/main/java/org/apache/samza/task/DrainListenerTask.java
@@ -0,0 +1,36 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.samza.task;
+
+/**
+ * The DrainListenerTask augments {@link StreamTask} allowing the method implementor to specify code to be
+ * executed when the 'drain' is reached for a task.
+ */
+public interface DrainListenerTask {
+ /**
+ * Guaranteed to be invoked when all SSPs processed by this task have drained.
+ *
+ * @param collector Contains the means of sending message envelopes to an output stream.*
+ * @param coordinator Manages execution of tasks.
+ *
+ * @throws Exception Any exception types encountered during the execution of the processing task.
+ */
+ void onDrain(MessageCollector collector, TaskCoordinator coordinator) throws Exception;
+}
diff --git a/samza-core/src/main/java/org/apache/samza/config/JobConfig.java b/samza-core/src/main/java/org/apache/samza/config/JobConfig.java
index 145f4cd..d36ac66 100644
--- a/samza-core/src/main/java/org/apache/samza/config/JobConfig.java
+++ b/samza-core/src/main/java/org/apache/samza/config/JobConfig.java
@@ -166,9 +166,12 @@
// Enable DrainMonitor in Samza Containers
// Default is false for now. Will be turned on after testing
- public static final String DRAIN_MONITOR_ENABLED = "samza.drain-monitor.enabled";
+ public static final String DRAIN_MONITOR_ENABLED = "job.drain-monitor.enabled";
public static final boolean DRAIN_MONITOR_ENABLED_DEFAULT = false;
+ public static final String DRAIN_MONITOR_POLL_INTERVAL_MILLIS = "job.drain-monitor.poll.interval.ms";
+ public static final long DRAIN_MONITOR_POLL_INTERVAL_MILLIS_DEFAULT = 60_000;
+
// Enable ClusterBasedJobCoordinator aka ApplicationMaster High Availability (AM-HA).
// High availability allows new AM to establish connection with already running containers
public static final String YARN_AM_HIGH_AVAILABILITY_ENABLED = "yarn.am.high-availability.enabled";
@@ -479,6 +482,10 @@
return getBoolean(DRAIN_MONITOR_ENABLED, DRAIN_MONITOR_ENABLED_DEFAULT);
}
+ public long getDrainMonitorPollIntervalMillis() {
+ return getLong(DRAIN_MONITOR_POLL_INTERVAL_MILLIS, DRAIN_MONITOR_POLL_INTERVAL_MILLIS_DEFAULT);
+ }
+
public long getContainerHeartbeatRetryCount() {
return getLong(YARN_CONTAINER_HEARTBEAT_RETRY_COUNT, YARN_CONTAINER_HEARTBEAT_RETRY_COUNT_DEFAULT);
}
diff --git a/samza-core/src/main/java/org/apache/samza/config/StreamConfig.java b/samza-core/src/main/java/org/apache/samza/config/StreamConfig.java
index 9f774b8..9fc74a7 100644
--- a/samza-core/src/main/java/org/apache/samza/config/StreamConfig.java
+++ b/samza-core/src/main/java/org/apache/samza/config/StreamConfig.java
@@ -167,7 +167,7 @@
* @param systemStream system stream to map to stream id
* @return stream id corresponding to the system stream
*/
- private String systemStreamToStreamId(SystemStream systemStream) {
+ public String systemStreamToStreamId(SystemStream systemStream) {
List<String> streamIds = getStreamIdsForSystem(systemStream.getSystem()).stream()
.filter(streamId -> systemStream.getStream().equals(getPhysicalName(streamId))).collect(Collectors.toList());
if (streamIds.size() > 1) {
diff --git a/samza-core/src/main/java/org/apache/samza/container/RunLoop.java b/samza-core/src/main/java/org/apache/samza/container/RunLoop.java
index b7a80a3..5bbcbae 100644
--- a/samza-core/src/main/java/org/apache/samza/container/RunLoop.java
+++ b/samza-core/src/main/java/org/apache/samza/container/RunLoop.java
@@ -24,6 +24,7 @@
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
+import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Optional;
@@ -35,6 +36,7 @@
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;
import org.apache.samza.SamzaException;
+import org.apache.samza.system.DrainMessage;
import org.apache.samza.system.IncomingMessageEnvelope;
import org.apache.samza.system.MessageType;
import org.apache.samza.system.SystemConsumers;
@@ -70,7 +72,6 @@
private final List<AsyncTaskWorker> taskWorkers;
private final SystemConsumers consumerMultiplexer;
private final Map<SystemStreamPartition, List<AsyncTaskWorker>> sspToTaskWorkerMapping;
-
private final ExecutorService threadPool;
private final CoordinatorRequests coordinatorRequests;
private final Object latch;
@@ -86,9 +87,11 @@
private volatile boolean shutdownNow = false;
private volatile Throwable throwable = null;
private final HighResolutionClock clock;
- private final boolean isAsyncCommitEnabled;
+ private boolean isAsyncCommitEnabled;
private volatile boolean runLoopResumedSinceLastChecked;
private final int elasticityFactor;
+ private final String runId;
+ private boolean isDraining = false;
public RunLoop(Map<TaskName, RunLoopTask> runLoopTasks,
ExecutorService threadPool,
@@ -103,7 +106,7 @@
HighResolutionClock clock,
boolean isAsyncCommitEnabled) {
this(runLoopTasks, threadPool, consumerMultiplexer, maxConcurrency, windowMs, commitMs, callbackTimeoutMs,
- maxThrottlingDelayMs, maxIdleMs, containerMetrics, clock, isAsyncCommitEnabled, 1);
+ maxThrottlingDelayMs, maxIdleMs, containerMetrics, clock, isAsyncCommitEnabled, 1, null);
}
public RunLoop(Map<TaskName, RunLoopTask> runLoopTasks,
@@ -118,7 +121,8 @@
SamzaContainerMetrics containerMetrics,
HighResolutionClock clock,
boolean isAsyncCommitEnabled,
- int elasticityFactor) {
+ int elasticityFactor,
+ String runId) {
this.threadPool = threadPool;
this.consumerMultiplexer = consumerMultiplexer;
@@ -134,18 +138,30 @@
this.latch = new Object();
this.workerTimer = Executors.newSingleThreadScheduledExecutor();
this.clock = clock;
- Map<TaskName, AsyncTaskWorker> workers = new HashMap<>();
+ // assign runId before creating workers. As the inner AsyncTaskWorker class is not static, it relies on
+ // the outer class fields to be init first
+ this.runId = runId;
+ Map<TaskName, AsyncTaskWorker> workers = new HashMap<>();
for (RunLoopTask task : runLoopTasks.values()) {
workers.put(task.taskName(), new AsyncTaskWorker(task));
}
// Partions and tasks assigned to the container will not change during the run loop life time
this.sspToTaskWorkerMapping = Collections.unmodifiableMap(getSspToAsyncTaskWorkerMap(runLoopTasks, workers));
+
this.taskWorkers = Collections.unmodifiableList(new ArrayList<>(workers.values()));
this.isAsyncCommitEnabled = isAsyncCommitEnabled;
this.elasticityFactor = elasticityFactor;
}
/**
+ * Sets the RunLoop to drain mode.
+ * */
+ private void drain() {
+ isDraining = true;
+ isAsyncCommitEnabled = false;
+ }
+
+ /**
* Returns mapping of the SystemStreamPartition to the AsyncTaskWorkers to efficiently route the envelopes
*/
private static Map<SystemStreamPartition, List<AsyncTaskWorker>> getSspToAsyncTaskWorkerMap(
@@ -297,7 +313,8 @@
* when elasticity is enabled,
* sspToTaskWorkerMapping has workers for a SSP which has keyBucket
* hence need to use envelop.getSSP(elasticityFactor)
- * Additionally, when envelope is EnofStream or Watermark, it needs to be sent to all works for the ssp irrespective of keyBucket
+ * Additionally, when envelope is EndOfStream or Watermark or Drain, it needs to be sent to all works for the ssp
+ * irrespective of keyBucket
* @param envelope
* @return list of workers for the envelope
*/
@@ -309,9 +326,14 @@
final SystemStreamPartition sspOfEnvelope = envelope.getSystemStreamPartition(elasticityFactor);
List<AsyncTaskWorker> listOfWorkersForEnvelope = null;
- // if envelope is end of stream or watermark, it needs to be routed to all tasks consuming the ssp irresp of keybucket
+ // if envelope is end of stream or watermark or drain, it needs to be routed to all tasks consuming the ssp irresp
+ // of keybucket
MessageType messageType = MessageType.of(envelope.getMessage());
- if (envelope.isEndOfStream() || MessageType.END_OF_STREAM == messageType || MessageType.WATERMARK == messageType) {
+ if (envelope.isEndOfStream()
+ || envelope.isDrain()
+ || messageType == MessageType.END_OF_STREAM
+ || messageType == MessageType.DRAIN
+ || messageType == MessageType.WATERMARK) {
//sspToTaskWorkerMapping has ssps with keybucket so extract and check only system, stream and partition and ignore the keybucket
listOfWorkersForEnvelope = sspToTaskWorkerMapping.entrySet()
@@ -391,6 +413,20 @@
}
/**
+ * Resume the runloop thread. This method is triggered after a task has completed drain.
+ */
+ private void resumeAfterDrain() {
+ log.trace("Resume loop thread");
+ if (coordinatorRequests.shouldShutdownNow()) {
+ shutdownNow = true;
+ }
+ synchronized (latch) {
+ latch.notifyAll();
+ runLoopResumedSinceLastChecked = true;
+ }
+ }
+
+ /**
* Set the throwable and abort run loop. The throwable will be thrown from the run loop thread
* @param t throwable
*/
@@ -426,6 +462,7 @@
COMMIT,
PROCESS,
END_OF_STREAM,
+ DRAIN,
SCHEDULER,
NO_OP
}
@@ -443,7 +480,8 @@
this.task = task;
this.callbackManager = new TaskCallbackManager(this, callbackTimer, callbackTimeoutMs, maxConcurrency, clock);
Set<SystemStreamPartition> sspSet = getWorkingSSPSet(task);
- this.state = new AsyncTaskState(task.taskName(), task.metrics(), sspSet, !task.intermediateStreams().isEmpty());
+ this.state = new AsyncTaskState(task.taskName(), task.metrics(), sspSet, !task.intermediateStreams().isEmpty(),
+ runId);
}
private void init() {
@@ -514,18 +552,46 @@
case END_OF_STREAM:
endOfStream();
break;
+ case DRAIN:
+ drain();
+ break;
default:
//no op
break;
}
}
+ /**
+ * Called when a task has drained i.e all SSPs for the task have received a drain message.
+ * */
+ private void drain() {
+ state.complete = true;
+ state.startDrain();
+ try {
+ ReadableCoordinator coordinator = new ReadableCoordinator(task.taskName());
+
+ task.drain(coordinator);
+
+ // issue a shutdown request for the task
+ coordinator.shutdown(TaskCoordinator.RequestScope.CURRENT_TASK);
+ coordinatorRequests.update(coordinator);
+
+ // issue a commit explicitly before we shutdown the task
+ // Adding commit to coordinator will not work as the state is marked complete and NO_OP will always be the
+ // next operation for this task
+ commit();
+ } finally {
+ resumeAfterDrain();
+ }
+ }
+
private void endOfStream() {
state.complete = true;
try {
ReadableCoordinator coordinator = new ReadableCoordinator(task.taskName());
task.endOfStream(coordinator);
+
// issue a request for shutdown of the task
coordinator.shutdown(TaskCoordinator.RequestScope.CURRENT_TASK);
coordinatorRequests.update(coordinator);
@@ -538,7 +604,6 @@
} finally {
resume();
}
-
}
/**
@@ -638,9 +703,12 @@
}
};
- if (threadPool != null) {
+ if (threadPool != null && !isDraining) {
log.trace("Task {} commits on the thread pool", task.taskName());
threadPool.submit(commitWorker);
+ } else if (isDraining) {
+ log.trace("Task {} commits on the run loop thread as task is draining", task.taskName());
+ commitWorker.run();
} else {
log.trace("Task {} commits on the run loop thread", task.taskName());
commitWorker.run();
@@ -759,24 +827,31 @@
private volatile boolean needScheduler = false;
private volatile boolean complete = false;
private volatile boolean endOfStream = false;
+ private volatile boolean shouldDrain = false;
private volatile boolean windowInFlight = false;
private volatile boolean commitInFlight = false;
private volatile boolean schedulerInFlight = false;
private final AtomicInteger messagesInFlight = new AtomicInteger(0);
- private final ArrayDeque<PendingEnvelope> pendingEnvelopeQueue;
+ private final ArrayDeque<PendingEnvelope> pendingEnvelopeQueue;
//Set of SSPs that we are currently processing for this task instance
private final Set<SystemStreamPartition> processingSspSet;
+ //Set of SSPs that we are currently processing for this task instance
+ private final Set<SystemStreamPartition> processingSspSetToDrain;
private final TaskName taskName;
private final TaskInstanceMetrics taskMetrics;
private final boolean hasIntermediateStreams;
+ private final String runId;
- AsyncTaskState(TaskName taskName, TaskInstanceMetrics taskMetrics, Set<SystemStreamPartition> sspSet, boolean hasIntermediateStreams) {
+ AsyncTaskState(TaskName taskName, TaskInstanceMetrics taskMetrics, Set<SystemStreamPartition> sspSet,
+ boolean hasIntermediateStreams, String runId) {
this.taskName = taskName;
this.taskMetrics = taskMetrics;
this.pendingEnvelopeQueue = new ArrayDeque<>();
this.processingSspSet = sspSet;
+ this.processingSspSetToDrain = new HashSet<>(sspSet);
this.hasIntermediateStreams = hasIntermediateStreams;
+ this.runId = runId;
}
private boolean checkEndOfStream() {
@@ -796,6 +871,7 @@
&& sspInSet.getPartition().equals(sspOfEnvelope.getPartition()))
.findFirst();
ssp.ifPresent(processingSspSet::remove);
+ ssp.ifPresent(processingSspSetToDrain::remove);
}
if (!hasIntermediateStreams) {
pendingEnvelopeQueue.remove();
@@ -805,6 +881,57 @@
return processingSspSet.isEmpty();
}
+ private boolean shouldDrain() {
+ if (endOfStream) {
+ return false;
+ }
+
+ if (!pendingEnvelopeQueue.isEmpty()) {
+ PendingEnvelope pendingEnvelope = pendingEnvelopeQueue.peek();
+ IncomingMessageEnvelope envelope = pendingEnvelope.envelope;
+
+ if (envelope.isDrain()) {
+ final DrainMessage message = (DrainMessage) envelope.getMessage();
+ if (!message.getRunId().equals(runId)) {
+ // Removing the drain message from the pending queue as it doesn't match with the current runId
+ // Removing it will ensure that it is not picked up by process()
+ pendingEnvelopeQueue.remove();
+ } else {
+ // set the RunLoop to drain mode
+ if (!isDraining) {
+ drain();
+ }
+
+ if (elasticityFactor <= 1) {
+ SystemStreamPartition ssp = envelope.getSystemStreamPartition();
+ processingSspSetToDrain.remove(ssp);
+ } else {
+ // SystemConsumers will write only one envelope (enclosing DrainMessage) per SSP in its buffer.
+ // This envelope doesn't have keybucket info it's SSP. With elasticity, the same SSP can be processed by
+ // multiple tasks. Therefore, if envelope contains drain message, the ssp of envelope should be removed
+ // from task's processing set irrespective of keyBucket.
+ SystemStreamPartition sspOfEnvelope = envelope.getSystemStreamPartition();
+ Optional<SystemStreamPartition> ssp = processingSspSetToDrain.stream()
+ .filter(sspInSet -> sspInSet.getSystemStream().equals(sspOfEnvelope.getSystemStream())
+ && sspInSet.getPartition().equals(sspOfEnvelope.getPartition()))
+ .findFirst();
+ ssp.ifPresent(processingSspSetToDrain::remove);
+ }
+
+ if (!hasIntermediateStreams) {
+ // Don't remove from the pending queue as we want the DAG to pick up Drain message and propagate it to
+ // intermediate streams
+ pendingEnvelopeQueue.remove();
+ }
+ }
+ }
+ return processingSspSetToDrain.isEmpty();
+ }
+ // if no messages are in the queue, the task has probably already drained or there are no messages from
+ // the chooser
+ return false;
+ }
+
/**
* Returns whether the task is ready to do process/window/commit.
*
@@ -813,6 +940,11 @@
if (checkEndOfStream()) {
endOfStream = true;
}
+
+ if (shouldDrain()) {
+ shouldDrain = true;
+ }
+
if (coordinatorRequests.commitRequests().remove(taskName)) {
needCommit = true;
}
@@ -826,9 +958,9 @@
*/
if (needCommit) {
return (messagesInFlight.get() == 0 || isAsyncCommitEnabled) && !opInFlight;
- } else if (needWindow || needScheduler || endOfStream) {
+ } else if (needWindow || needScheduler || endOfStream || shouldDrain) {
/*
- * A task is ready for window, scheduler or end-of-stream operation.
+ * A task is ready for window, scheduler, drain or end-of-stream operation.
*/
return messagesInFlight.get() == 0 && !opInFlight;
} else {
@@ -847,14 +979,24 @@
*/
private WorkerOp nextOp() {
- if (complete) return WorkerOp.NO_OP;
+ if (complete) {
+ return WorkerOp.NO_OP;
+ }
if (isReady()) {
- if (needCommit) return WorkerOp.COMMIT;
- else if (needWindow) return WorkerOp.WINDOW;
- else if (needScheduler) return WorkerOp.SCHEDULER;
- else if (endOfStream && pendingEnvelopeQueue.isEmpty()) return WorkerOp.END_OF_STREAM;
- else if (!pendingEnvelopeQueue.isEmpty()) return WorkerOp.PROCESS;
+ if (needCommit) {
+ return WorkerOp.COMMIT;
+ } else if (needWindow) {
+ return WorkerOp.WINDOW;
+ } else if (needScheduler) {
+ return WorkerOp.SCHEDULER;
+ } else if (endOfStream && pendingEnvelopeQueue.isEmpty()) {
+ return WorkerOp.END_OF_STREAM;
+ } else if (shouldDrain && pendingEnvelopeQueue.isEmpty()) {
+ return WorkerOp.DRAIN;
+ } else if (!pendingEnvelopeQueue.isEmpty()) {
+ return WorkerOp.PROCESS;
+ }
}
return WorkerOp.NO_OP;
}
@@ -876,6 +1018,10 @@
windowInFlight = true;
}
+ private void startDrain() {
+ shouldDrain = false;
+ }
+
private void startCommit() {
needCommit = false;
commitInFlight = true;
@@ -918,6 +1064,7 @@
int queueSize = pendingEnvelopeQueue.size();
taskMetrics.pendingMessages().set(queueSize);
log.trace("Insert envelope to task {} queue.", taskName);
+ log.trace("Insert envelope to task {} queue.", taskName);
log.debug("Task {} pending envelope count is {} after insertion.", taskName, queueSize);
}
diff --git a/samza-core/src/main/java/org/apache/samza/container/RunLoopFactory.java b/samza-core/src/main/java/org/apache/samza/container/RunLoopFactory.java
index 9d20950..cb021d0 100644
--- a/samza-core/src/main/java/org/apache/samza/container/RunLoopFactory.java
+++ b/samza-core/src/main/java/org/apache/samza/container/RunLoopFactory.java
@@ -41,7 +41,8 @@
SamzaContainerMetrics containerMetrics,
TaskConfig taskConfig,
HighResolutionClock clock,
- int elasticityFactor) {
+ int elasticityFactor,
+ String runId) {
long taskWindowMs = taskConfig.getWindowMs();
@@ -65,6 +66,8 @@
log.info("Got elasticity factor: {}.", elasticityFactor);
+ log.info("Got current run Id: {}.", runId);
+
log.info("Run loop in asynchronous mode.");
return new RunLoop(
@@ -80,6 +83,7 @@
containerMetrics,
clock,
isAsyncCommitEnabled,
- elasticityFactor);
+ elasticityFactor,
+ runId);
}
}
diff --git a/samza-core/src/main/java/org/apache/samza/container/RunLoopTask.java b/samza-core/src/main/java/org/apache/samza/container/RunLoopTask.java
index 551da88..af63494 100644
--- a/samza-core/src/main/java/org/apache/samza/container/RunLoopTask.java
+++ b/samza-core/src/main/java/org/apache/samza/container/RunLoopTask.java
@@ -99,6 +99,15 @@
void endOfStream(ReadableCoordinator coordinator);
/**
+ * Called when all {@link SystemStreamPartition} processed by a task have drained. This is called only
+ * once per task. {@link RunLoop} will issue a shutdown request to the coordinator immediately following the
+ * invocation of this method.
+ *
+ * @param coordinator manages execution of tasks.
+ */
+ void drain(ReadableCoordinator coordinator);
+
+ /**
* Indicates whether {@link #window} should be invoked on this task. If true, {@link RunLoop}
* will schedule window to execute periodically according to its windowMs.
*
diff --git a/samza-core/src/main/java/org/apache/samza/drain/DrainMonitor.java b/samza-core/src/main/java/org/apache/samza/drain/DrainMonitor.java
index 6b5c98e..b46fc3f 100644
--- a/samza-core/src/main/java/org/apache/samza/drain/DrainMonitor.java
+++ b/samza-core/src/main/java/org/apache/samza/drain/DrainMonitor.java
@@ -33,6 +33,7 @@
import org.apache.samza.SamzaException;
import org.apache.samza.config.ApplicationConfig;
import org.apache.samza.config.Config;
+import org.apache.samza.config.JobConfig;
import org.apache.samza.coordinator.metadatastore.NamespaceAwareCoordinatorStreamStore;
import org.apache.samza.metadatastore.MetadataStore;
import org.slf4j.Logger;
@@ -65,7 +66,6 @@
STOPPED
}
- private static final int POLLING_INTERVAL_MILLIS = 60_000;
private static final int INITIAL_POLL_DELAY_MILLIS = 0;
private final ScheduledExecutorService schedulerService =
@@ -75,7 +75,6 @@
.setDaemon(true)
.build());
private final String appRunId;
- private final long pollingIntervalMillis;
private final NamespaceAwareCoordinatorStreamStore drainMetadataStore;
// Used to guard write access to state.
private final Object lock = new Object();
@@ -83,20 +82,23 @@
@GuardedBy("lock")
private State state = State.INIT;
private DrainCallback callback;
+ private long pollingIntervalMillis;
public DrainMonitor(MetadataStore metadataStore, Config config) {
- this(metadataStore, config, POLLING_INTERVAL_MILLIS);
- }
-
- public DrainMonitor(MetadataStore metadataStore, Config config, long pollingIntervalMillis) {
Preconditions.checkNotNull(metadataStore, "MetadataStore parameter cannot be null.");
Preconditions.checkNotNull(config, "Config parameter cannot be null.");
- Preconditions.checkArgument(pollingIntervalMillis > 0,
- String.format("Polling interval specified is %d ms. It should be greater than 0.", pollingIntervalMillis));
this.drainMetadataStore =
new NamespaceAwareCoordinatorStreamStore(metadataStore, DrainUtils.DRAIN_METADATA_STORE_NAMESPACE);
ApplicationConfig applicationConfig = new ApplicationConfig(config);
this.appRunId = applicationConfig.getRunId();
+ JobConfig jobConfig = new JobConfig(config);
+ this.pollingIntervalMillis = jobConfig.getDrainMonitorPollIntervalMillis();
+ }
+
+ public DrainMonitor(MetadataStore metadataStore, Config config, long pollingIntervalMillis) {
+ this(metadataStore, config);
+ Preconditions.checkArgument(pollingIntervalMillis > 0,
+ String.format("Polling interval specified is %d ms. It should be greater than 0.", pollingIntervalMillis));
this.pollingIntervalMillis = pollingIntervalMillis;
}
@@ -207,12 +209,12 @@
* One time check check to see if there are any DrainNotification messages available in the
* metadata store for the current deployment.
* */
- static boolean shouldDrain(NamespaceAwareCoordinatorStreamStore drainMetadataStore, String deploymentId) {
+ static boolean shouldDrain(NamespaceAwareCoordinatorStreamStore drainMetadataStore, String runId) {
final Optional<List<DrainNotification>> drainNotifications = readDrainNotificationMessages(drainMetadataStore);
if (drainNotifications.isPresent()) {
final ImmutableList<DrainNotification> filteredDrainNotifications = drainNotifications.get()
.stream()
- .filter(notification -> deploymentId.equals(notification.getDeploymentId()))
+ .filter(notification -> runId.equals(notification.getRunId()))
.collect(ImmutableList.toImmutableList());
return !filteredDrainNotifications.isEmpty();
}
diff --git a/samza-core/src/main/java/org/apache/samza/drain/DrainNotification.java b/samza-core/src/main/java/org/apache/samza/drain/DrainNotification.java
index a16595e..dd97cd7 100644
--- a/samza-core/src/main/java/org/apache/samza/drain/DrainNotification.java
+++ b/samza-core/src/main/java/org/apache/samza/drain/DrainNotification.java
@@ -33,26 +33,26 @@
/**
* Unique identifier for a deployment so drain notifications messages can be invalidated across a job restarts.
*/
- private final String deploymentId;
+ private final String runId;
- public DrainNotification(UUID uuid, String deploymentId) {
+ public DrainNotification(UUID uuid, String runId) {
this.uuid = uuid;
- this.deploymentId = deploymentId;
+ this.runId = runId;
}
public UUID getUuid() {
return this.uuid;
}
- public String getDeploymentId() {
- return deploymentId;
+ public String getRunId() {
+ return runId;
}
@Override
public String toString() {
final StringBuilder sb = new StringBuilder("DrainMessage{");
sb.append(" UUID: ").append(uuid);
- sb.append(", deploymentId: '").append(deploymentId).append('\'');
+ sb.append(", runId: '").append(runId).append('\'');
sb.append('}');
return sb.toString();
}
@@ -67,11 +67,11 @@
}
DrainNotification that = (DrainNotification) o;
return Objects.equal(uuid, that.uuid)
- && Objects.equal(deploymentId, that.deploymentId);
+ && Objects.equal(runId, that.runId);
}
@Override
public int hashCode() {
- return Objects.hashCode(uuid, deploymentId);
+ return Objects.hashCode(uuid, runId);
}
}
diff --git a/samza-core/src/main/java/org/apache/samza/drain/DrainNotificationObjectMapper.java b/samza-core/src/main/java/org/apache/samza/drain/DrainNotificationObjectMapper.java
index 48e6db1..72cdfbb 100644
--- a/samza-core/src/main/java/org/apache/samza/drain/DrainNotificationObjectMapper.java
+++ b/samza-core/src/main/java/org/apache/samza/drain/DrainNotificationObjectMapper.java
@@ -67,7 +67,7 @@
throws IOException {
Map<String, Object> drainMessageMap = new HashMap<>();
drainMessageMap.put("uuid", value.getUuid().toString());
- drainMessageMap.put("deploymentId", value.getDeploymentId());
+ drainMessageMap.put("runId", value.getRunId());
jsonGenerator.writeObject(drainMessageMap);
}
}
@@ -79,8 +79,8 @@
ObjectCodec oc = jsonParser.getCodec();
JsonNode node = oc.readTree(jsonParser);
UUID uuid = UUID.fromString(node.get("uuid").textValue());
- String deploymentId = node.get("deploymentId").textValue();
- return new DrainNotification(uuid, deploymentId);
+ String runId = node.get("runId").textValue();
+ return new DrainNotification(uuid, runId);
}
}
}
diff --git a/samza-core/src/main/java/org/apache/samza/drain/DrainUtils.java b/samza-core/src/main/java/org/apache/samza/drain/DrainUtils.java
index c100a47..f712381 100644
--- a/samza-core/src/main/java/org/apache/samza/drain/DrainUtils.java
+++ b/samza-core/src/main/java/org/apache/samza/drain/DrainUtils.java
@@ -47,23 +47,23 @@
* Writes a {@link DrainNotification} to the underlying metastore. This method should be used by external controllers
* to issue a DrainNotification to the JobCoordinator and Samza Containers.
* @param metadataStore Metadata store to write drain notification to.
- * @param deploymentId deploymentId for the DrainNotification
+ * @param runId runId for the DrainNotification
*
* @return generated uuid for the DrainNotification
*/
- public static UUID writeDrainNotification(MetadataStore metadataStore, String deploymentId) {
+ public static UUID writeDrainNotification(MetadataStore metadataStore, String runId) {
Preconditions.checkArgument(metadataStore != null, "MetadataStore cannot be null.");
- Preconditions.checkArgument(!Strings.isNullOrEmpty(deploymentId), "deploymentId should be non-null.");
- LOG.info("Attempting to write DrainNotification to metadata-store for the deployment ID {}", deploymentId);
+ Preconditions.checkArgument(!Strings.isNullOrEmpty(runId), "runId should be non-null.");
+ LOG.info("Attempting to write DrainNotification to metadata-store for the deployment ID {}", runId);
final NamespaceAwareCoordinatorStreamStore drainMetadataStore =
new NamespaceAwareCoordinatorStreamStore(metadataStore, DRAIN_METADATA_STORE_NAMESPACE);
final ObjectMapper objectMapper = DrainNotificationObjectMapper.getObjectMapper();
final UUID uuid = UUID.randomUUID();
- final DrainNotification message = new DrainNotification(uuid, deploymentId);
+ final DrainNotification message = new DrainNotification(uuid, runId);
try {
drainMetadataStore.put(message.getUuid().toString(), objectMapper.writeValueAsBytes(message));
drainMetadataStore.flush();
- LOG.info("DrainNotification with id {} written to metadata-store for the deployment ID {}", uuid, deploymentId);
+ LOG.info("DrainNotification with id {} written to metadata-store for the deployment ID {}", uuid, runId);
} catch (Exception ex) {
throw new SamzaException(
String.format("DrainNotification might have been not written to metastore %s", message), ex);
@@ -73,23 +73,23 @@
/**
* Cleans up DrainNotifications for the current deployment from the underlying metadata store.
- * The current deploymentId is extracted from the config.
+ * The current runId is extracted from the config.
*
* @param metadataStore underlying metadata store
- * @param config Config for the job. Used to extract the deploymentId of the job.
+ * @param config Config for the job. Used to extract the runId of the job.
* */
public static void cleanup(MetadataStore metadataStore, Config config) {
Preconditions.checkArgument(metadataStore != null, "MetadataStore cannot be null.");
Preconditions.checkNotNull(config, "Config parameter cannot be null.");
final ApplicationConfig applicationConfig = new ApplicationConfig(config);
- final String deploymentId = applicationConfig.getRunId();
+ final String runId = applicationConfig.getRunId();
final ObjectMapper objectMapper = DrainNotificationObjectMapper.getObjectMapper();
final NamespaceAwareCoordinatorStreamStore drainMetadataStore =
new NamespaceAwareCoordinatorStreamStore(metadataStore, DRAIN_METADATA_STORE_NAMESPACE);
- if (DrainMonitor.shouldDrain(drainMetadataStore, deploymentId)) {
- LOG.info("Attempting to clean up DrainNotifications from the metadata-store for the current deployment {}", deploymentId);
+ if (DrainMonitor.shouldDrain(drainMetadataStore, runId)) {
+ LOG.info("Attempting to clean up DrainNotifications from the metadata-store for the current deployment {}", runId);
drainMetadataStore.all()
.values()
.stream()
@@ -101,19 +101,19 @@
throw new SamzaException(e);
}
})
- .filter(notification -> deploymentId.equals(notification.getDeploymentId()))
+ .filter(notification -> runId.equals(notification.getRunId()))
.forEach(notification -> drainMetadataStore.delete(notification.getUuid().toString()));
drainMetadataStore.flush();
- LOG.info("Successfully cleaned up DrainNotifications from the metadata-store for the current deployment {}", deploymentId);
+ LOG.info("Successfully cleaned up DrainNotifications from the metadata-store for the current deployment {}", runId);
} else {
LOG.info("No DrainNotification found in the metadata-store for the current deployment {}. No need to cleanup.",
- deploymentId);
+ runId);
}
}
/**
- * Cleans up all DrainNotifications irrespective of the deploymentId.
+ * Cleans up all DrainNotifications irrespective of the runId.
* */
public static void cleanupAll(MetadataStore metadataStore) {
Preconditions.checkArgument(metadataStore != null, "MetadataStore cannot be null.");
diff --git a/samza-core/src/main/java/org/apache/samza/operators/impl/BroadcastOperatorImpl.java b/samza-core/src/main/java/org/apache/samza/operators/impl/BroadcastOperatorImpl.java
index 4f93f2c..7dd5d19 100644
--- a/samza-core/src/main/java/org/apache/samza/operators/impl/BroadcastOperatorImpl.java
+++ b/samza-core/src/main/java/org/apache/samza/operators/impl/BroadcastOperatorImpl.java
@@ -21,10 +21,12 @@
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionStage;
+import org.apache.samza.config.ApplicationConfig;
import org.apache.samza.context.Context;
import org.apache.samza.operators.spec.BroadcastOperatorSpec;
import org.apache.samza.operators.spec.OperatorSpec;
import org.apache.samza.system.ControlMessage;
+import org.apache.samza.system.DrainMessage;
import org.apache.samza.system.EndOfStreamMessage;
import org.apache.samza.system.OutgoingMessageEnvelope;
import org.apache.samza.system.SystemStream;
@@ -40,11 +42,13 @@
private final BroadcastOperatorSpec<M> broadcastOpSpec;
private final SystemStream systemStream;
private final String taskName;
+ private final String runId;
BroadcastOperatorImpl(BroadcastOperatorSpec<M> broadcastOpSpec, SystemStream systemStream, Context context) {
this.broadcastOpSpec = broadcastOpSpec;
this.systemStream = systemStream;
this.taskName = context.getTaskContext().getTaskModel().getTaskName().getTaskName();
+ this.runId = new ApplicationConfig(context.getJobContext().getConfig()).getRunId();
}
@Override
@@ -74,6 +78,12 @@
}
@Override
+ protected Collection<Void> handleDrain(MessageCollector collector, TaskCoordinator coordinator) {
+ sendControlMessage(new DrainMessage(taskName, runId), collector);
+ return Collections.emptyList();
+ }
+
+ @Override
protected Collection<Void> handleWatermark(long watermark, MessageCollector collector, TaskCoordinator coordinator) {
sendControlMessage(new WatermarkMessage(watermark, taskName), collector);
return Collections.emptyList();
diff --git a/samza-core/src/main/java/org/apache/samza/operators/impl/DrainStates.java b/samza-core/src/main/java/org/apache/samza/operators/impl/DrainStates.java
new file mode 100644
index 0000000..30a341f
--- /dev/null
+++ b/samza-core/src/main/java/org/apache/samza/operators/impl/DrainStates.java
@@ -0,0 +1,119 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.samza.operators.impl;
+
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
+import java.util.stream.Collectors;
+import org.apache.samza.system.DrainMessage;
+import org.apache.samza.system.SystemStream;
+import org.apache.samza.system.SystemStreamPartition;
+
+
+/**
+ * This class tracks the drain state for streams in a task. Internally it keeps track of Drain messages received
+ * from upstream tasks for each system stream partition (ssp). If messages have been received from all tasks,
+ * it will mark the ssp as drained. For a stream to be drained, all its partitions assigned to
+ * the task need to be drained.
+ *
+ * This class is thread-safe.
+ */
+public class DrainStates {
+ private static final class DrainState {
+ // set of upstream tasks
+ private final Set<String> tasks = new HashSet<>();
+ private final int expectedTotal;
+ private volatile boolean drained = false;
+
+ DrainState(int expectedTotal) {
+ this.expectedTotal = expectedTotal;
+ }
+
+ synchronized void update(String taskName) {
+ if (taskName != null) {
+ // aggregate the eos messages
+ tasks.add(taskName);
+ drained = tasks.size() == expectedTotal;
+ } else {
+ // eos is coming from either source or aggregator task
+ drained = true;
+ }
+ }
+
+ boolean isDrained() {
+ return drained;
+ }
+
+ @Override
+ public String toString() {
+ return "DrainState: [Tasks : "
+ + tasks
+ + ", isDrained : "
+ + drained
+ + "]";
+ }
+ }
+
+ private final Map<SystemStreamPartition, DrainState> drainStates;
+
+ /**
+ * Constructing the drain states for a task.
+ * @param ssps all the ssps assigned to this task
+ * @param producerTaskCounts mapping from a stream to the number of upstream tasks that produce to it
+ */
+ DrainStates(Set<SystemStreamPartition> ssps, Map<SystemStream, Integer> producerTaskCounts) {
+ this.drainStates = ssps.stream()
+ .collect(Collectors.toMap(
+ ssp -> ssp,
+ ssp -> new DrainState(producerTaskCounts.getOrDefault(ssp.getSystemStream(), 0))));
+ }
+
+ /**
+ * Update the state upon receiving a drain message.
+ * @param eos message of {@link DrainMessage}
+ * @param ssp system stream partition
+ */
+ void update(DrainMessage eos, SystemStreamPartition ssp) {
+ DrainState state = drainStates.get(ssp);
+ state.update(eos.getTaskName());
+ }
+
+ /**
+ * Checks if the system-stream is drained.
+ * */
+ boolean isDrained(SystemStream systemStream) {
+ return drainStates.entrySet().stream()
+ .filter(entry -> entry.getKey().getSystemStream().equals(systemStream))
+ .allMatch(entry -> entry.getValue().isDrained());
+ }
+
+ /**
+ * Checks if all streams (input SSPs) for the task has drained.
+ * */
+ boolean areAllStreamsDrained() {
+ return drainStates.values().stream().allMatch(DrainState::isDrained);
+ }
+
+ @Override
+ public String toString() {
+ return drainStates.toString();
+ }
+}
diff --git a/samza-core/src/main/java/org/apache/samza/operators/impl/OperatorImpl.java b/samza-core/src/main/java/org/apache/samza/operators/impl/OperatorImpl.java
index c404475..5723d91 100644
--- a/samza-core/src/main/java/org/apache/samza/operators/impl/OperatorImpl.java
+++ b/samza-core/src/main/java/org/apache/samza/operators/impl/OperatorImpl.java
@@ -1,4 +1,4 @@
-/*
+ /*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
@@ -39,6 +39,7 @@
import org.apache.samza.operators.functions.WatermarkFunction;
import org.apache.samza.operators.spec.OperatorSpec;
import org.apache.samza.scheduler.CallbackScheduler;
+import org.apache.samza.system.DrainMessage;
import org.apache.samza.system.EndOfStreamMessage;
import org.apache.samza.system.SystemStream;
import org.apache.samza.system.SystemStreamPartition;
@@ -86,6 +87,8 @@
private TaskModel taskModel;
// end-of-stream states
private EndOfStreamStates eosStates;
+ // drain states
+ private DrainStates drainStates;
// watermark states
private WatermarkStates watermarkStates;
private CallbackScheduler callbackScheduler;
@@ -125,6 +128,7 @@
this.taskName = taskContext.getTaskModel().getTaskName();
this.eosStates = (EndOfStreamStates) internalTaskContext.fetchObject(EndOfStreamStates.class.getName());
this.watermarkStates = (WatermarkStates) internalTaskContext.fetchObject(WatermarkStates.class.getName());
+ this.drainStates = (DrainStates) internalTaskContext.fetchObject(DrainStates.class.getName());
this.controlMessageSender = new ControlMessageSender(internalTaskContext.getStreamMetadataCache());
this.taskModel = taskContext.getTaskModel();
this.callbackScheduler = taskContext.getCallbackScheduler();
@@ -362,6 +366,88 @@
}
/**
+ * This method is implemented when all input stream to this operation have encountered drain-and-stop control message.
+ * Inherited class should handle drain-and-stop by overriding this function.
+ * By default noop implementation is for in-memory operator to handle the drain-and-stop. Output operator need to
+ * override this to actually propagate drain-and-stop over the wire.
+ * @param collector message collector
+ * @param coordinator task coordinator
+ * @return results to be emitted when this operator reaches drain-and-stop
+ */
+ protected Collection<RM> handleDrain(MessageCollector collector, TaskCoordinator coordinator) {
+ return Collections.emptyList();
+ }
+
+ /**
+ * Aggregate {@link DrainMessage} from each ssp of the stream.
+ * Invoke {@link #onDrainOfStream(MessageCollector, TaskCoordinator)} if the stream reaches the end.
+ * @param drainMessage {@link DrainMessage} object
+ * @param ssp system stream partition
+ * @param collector message collector
+ * @param coordinator task coordinator
+ */
+ public final CompletionStage<Void> aggregateDrainMessages(DrainMessage drainMessage, SystemStreamPartition ssp,
+ MessageCollector collector, TaskCoordinator coordinator) {
+ LOG.info("Received drain message from task {} in {}", drainMessage.getTaskName(), ssp);
+ drainStates.update(drainMessage, ssp);
+
+ SystemStream stream = ssp.getSystemStream();
+ CompletionStage<Void> drainFuture = CompletableFuture.completedFuture(null);
+
+ if (drainStates.isDrained(stream)) {
+ LOG.info("Input {} reaches the end for task {}", stream.toString(), taskName.getTaskName());
+ if (drainMessage.getTaskName() != null && shouldTaskBroadcastToOtherPartitions(ssp)) {
+ // This is the aggregation task, which already received all the eos messages from upstream
+ // broadcast the end-of-stream to all the peer partitions
+ // additionally if elasiticty is enabled
+ // then only one of the elastic tasks of the ssp will broadcast
+ controlMessageSender.broadcastToOtherPartitions(new DrainMessage(drainMessage.getRunId()), ssp, collector);
+ }
+
+ drainFuture = onDrainOfStream(collector, coordinator)
+ .thenAccept(result -> {
+ if (drainStates.areAllStreamsDrained()) {
+ // all input streams have been drained, shut down the task
+ LOG.info("All input streams have been drained for task {}", taskName.getTaskName());
+ coordinator.shutdown(TaskCoordinator.RequestScope.CURRENT_TASK);
+ }
+ });
+ }
+
+ return drainFuture;
+ }
+
+
+ /**
+ * Invoke {@link #handleDrain(MessageCollector, TaskCoordinator)} if all the input streams to the current operator
+ * have encountered drain message.
+ * Propagate the drain to downstream operators.
+ * @param collector message collector
+ * @param coordinator task coordinator
+ */
+ private CompletionStage<Void> onDrainOfStream(MessageCollector collector, TaskCoordinator coordinator) {
+ CompletionStage<Void> drainFuture = CompletableFuture.completedFuture(null);
+
+ if (inputStreams.stream().allMatch(input -> drainStates.isDrained(input))) {
+ Collection<RM> results = handleDrain(collector, coordinator);
+
+ CompletionStage<Void> resultFuture = CompletableFuture.allOf(
+ results.stream()
+ .flatMap(r -> this.registeredOperators.stream()
+ .map(op -> op.onMessageAsync(r, collector, coordinator)))
+ .toArray(CompletableFuture[]::new));
+
+ // propagate DrainMessage to downstream operators
+ drainFuture = resultFuture.thenCompose(x ->
+ CompletableFuture.allOf(this.registeredOperators.stream()
+ .map(op -> op.onDrainOfStream(collector, coordinator))
+ .toArray(CompletableFuture[]::new)));
+ }
+
+ return drainFuture;
+ }
+
+ /**
* Aggregate the {@link WatermarkMessage} from each ssp into a watermark. Then call onWatermark() if
* a new watermark exits.
* @param watermarkMessage a {@link WatermarkMessage} object
diff --git a/samza-core/src/main/java/org/apache/samza/operators/impl/OperatorImplGraph.java b/samza-core/src/main/java/org/apache/samza/operators/impl/OperatorImplGraph.java
index 1aa67d8..c62b0b2 100644
--- a/samza-core/src/main/java/org/apache/samza/operators/impl/OperatorImplGraph.java
+++ b/samza-core/src/main/java/org/apache/samza/operators/impl/OperatorImplGraph.java
@@ -115,11 +115,16 @@
// set states for end-of-stream; don't include side inputs (see SAMZA-2303)
internalTaskContext.registerObject(EndOfStreamStates.class.getName(),
new EndOfStreamStates(internalTaskContext.getSspsExcludingSideInputs(), producerTaskCounts));
+
// set states for watermark; don't include side inputs (see SAMZA-2303)
internalTaskContext.registerObject(WatermarkStates.class.getName(),
new WatermarkStates(internalTaskContext.getSspsExcludingSideInputs(), producerTaskCounts,
context.getContainerContext().getContainerMetricsRegistry()));
+ // set states for drain; don't include side inputs (see SAMZA-2303)
+ internalTaskContext.registerObject(DrainStates.class.getName(),
+ new DrainStates(internalTaskContext.getSspsExcludingSideInputs(), producerTaskCounts));
+
specGraph.getInputOperators().forEach((streamId, inputOpSpec) -> {
SystemStream systemStream = streamConfig.streamIdToSystemStream(streamId);
InputOperatorImpl inputOperatorImpl =
diff --git a/samza-core/src/main/java/org/apache/samza/operators/impl/PartitionByOperatorImpl.java b/samza-core/src/main/java/org/apache/samza/operators/impl/PartitionByOperatorImpl.java
index 47ad4f6..643abdf 100644
--- a/samza-core/src/main/java/org/apache/samza/operators/impl/PartitionByOperatorImpl.java
+++ b/samza-core/src/main/java/org/apache/samza/operators/impl/PartitionByOperatorImpl.java
@@ -20,12 +20,14 @@
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionStage;
+import org.apache.samza.config.ApplicationConfig;
import org.apache.samza.context.Context;
import org.apache.samza.context.InternalTaskContext;
import org.apache.samza.operators.functions.MapFunction;
import org.apache.samza.operators.spec.OperatorSpec;
import org.apache.samza.operators.spec.PartitionByOperatorSpec;
import org.apache.samza.system.ControlMessage;
+import org.apache.samza.system.DrainMessage;
import org.apache.samza.system.EndOfStreamMessage;
import org.apache.samza.system.OutgoingMessageEnvelope;
import org.apache.samza.system.StreamMetadataCache;
@@ -48,6 +50,7 @@
private final MapFunction<? super M, ? extends K> keyFunction;
private final MapFunction<? super M, ? extends V> valueFunction;
private final String taskName;
+ private final String runId;
private final ControlMessageSender controlMessageSender;
PartitionByOperatorImpl(PartitionByOperatorSpec<M, K, V> partitionByOpSpec,
@@ -57,6 +60,7 @@
this.keyFunction = partitionByOpSpec.getKeyFunction();
this.valueFunction = partitionByOpSpec.getValueFunction();
this.taskName = internalTaskContext.getContext().getTaskContext().getTaskModel().getTaskName().getTaskName();
+ this.runId = new ApplicationConfig(internalTaskContext.getContext().getJobContext().getConfig()).getRunId();
StreamMetadataCache streamMetadataCache = internalTaskContext.getStreamMetadataCache();
this.controlMessageSender = new ControlMessageSender(streamMetadataCache);
}
@@ -95,6 +99,12 @@
}
@Override
+ protected Collection<Void> handleDrain(MessageCollector collector, TaskCoordinator coordinator) {
+ sendControlMessage(new DrainMessage(taskName, runId), collector);
+ return Collections.emptyList();
+ }
+
+ @Override
protected Collection<Void> handleWatermark(long watermark, MessageCollector collector, TaskCoordinator coordinator) {
sendControlMessage(new WatermarkMessage(watermark, taskName), collector);
return Collections.emptyList();
diff --git a/samza-core/src/main/java/org/apache/samza/processor/StreamProcessor.java b/samza-core/src/main/java/org/apache/samza/processor/StreamProcessor.java
index 654d9a0..cda04377 100644
--- a/samza-core/src/main/java/org/apache/samza/processor/StreamProcessor.java
+++ b/samza-core/src/main/java/org/apache/samza/processor/StreamProcessor.java
@@ -48,6 +48,7 @@
import org.apache.samza.coordinator.JobCoordinatorFactory;
import org.apache.samza.coordinator.JobCoordinatorListener;
import org.apache.samza.diagnostics.DiagnosticsManager;
+import org.apache.samza.drain.DrainMonitor;
import org.apache.samza.job.model.JobModel;
import org.apache.samza.metadatastore.MetadataStore;
import org.apache.samza.metrics.MetricsRegistry;
@@ -402,12 +403,18 @@
*/
MetricsRegistryMap metricsRegistryMap = new MetricsRegistryMap();
+ DrainMonitor drainMonitor = null;
+ JobConfig jobConfig = new JobConfig(config);
+ if (metadataStore != null && jobConfig.getDrainMonitorEnabled()) {
+ drainMonitor = new DrainMonitor(metadataStore, config);
+ }
+
return SamzaContainer.apply(processorId, jobModel, ScalaJavaUtil.toScalaMap(this.customMetricsReporter),
metricsRegistryMap, this.taskFactory, JobContextImpl.fromConfigWithDefaults(this.config, jobModel),
Option.apply(this.applicationDefinedContainerContextFactoryOptional.orElse(null)),
Option.apply(this.applicationDefinedTaskContextFactoryOptional.orElse(null)),
Option.apply(this.externalContextOptional.orElse(null)), null, startpointManager,
- Option.apply(diagnosticsManager.orElse(null)), null);
+ Option.apply(diagnosticsManager.orElse(null)), drainMonitor);
}
private static JobCoordinator createJobCoordinator(Config config, String processorId, MetricsRegistry metricsRegistry, MetadataStore metadataStore) {
diff --git a/samza-core/src/main/java/org/apache/samza/runtime/ContainerLaunchUtil.java b/samza-core/src/main/java/org/apache/samza/runtime/ContainerLaunchUtil.java
index f499eb3..cd153b5 100644
--- a/samza-core/src/main/java/org/apache/samza/runtime/ContainerLaunchUtil.java
+++ b/samza-core/src/main/java/org/apache/samza/runtime/ContainerLaunchUtil.java
@@ -143,7 +143,8 @@
MetricsRegistryMap metricsRegistryMap = new MetricsRegistryMap();
DrainMonitor drainMonitor = null;
- if (new JobConfig(config).getDrainMonitorEnabled()) {
+ JobConfig jobConfig = new JobConfig(config);
+ if (jobConfig.getDrainMonitorEnabled()) {
drainMonitor = new DrainMonitor(coordinatorStreamStore, config);
}
diff --git a/samza-core/src/main/java/org/apache/samza/serializers/IntermediateMessageSerde.java b/samza-core/src/main/java/org/apache/samza/serializers/IntermediateMessageSerde.java
index 83a0a35..a7afecd 100644
--- a/samza-core/src/main/java/org/apache/samza/serializers/IntermediateMessageSerde.java
+++ b/samza-core/src/main/java/org/apache/samza/serializers/IntermediateMessageSerde.java
@@ -22,6 +22,7 @@
import java.util.Arrays;
import org.apache.samza.SamzaException;
+import org.apache.samza.system.DrainMessage;
import org.apache.samza.system.EndOfStreamMessage;
import org.apache.samza.system.MessageType;
import org.apache.samza.system.WatermarkMessage;
@@ -56,11 +57,13 @@
private final Serde userMessageSerde;
private final Serde<WatermarkMessage> watermarkSerde;
private final Serde<EndOfStreamMessage> eosSerde;
+ private final Serde<DrainMessage> drainMessageSerde;
public IntermediateMessageSerde(Serde userMessageSerde) {
this.userMessageSerde = userMessageSerde;
this.watermarkSerde = new JsonSerdeV2<>(WatermarkMessage.class);
this.eosSerde = new JsonSerdeV2<>(EndOfStreamMessage.class);
+ this.drainMessageSerde = new JsonSerdeV2<>(DrainMessage.class);
}
@Override
@@ -93,6 +96,9 @@
case END_OF_STREAM:
object = eosSerde.fromBytes(data);
break;
+ case DRAIN:
+ object = drainMessageSerde.fromBytes(data);
+ break;
default:
throw new UnsupportedOperationException(String.format("Message type %s is not supported", type.name()));
}
@@ -118,6 +124,9 @@
case END_OF_STREAM:
data = eosSerde.toBytes((EndOfStreamMessage) object);
break;
+ case DRAIN:
+ data = drainMessageSerde.toBytes((DrainMessage) object);
+ break;
default:
throw new SamzaException("Unknown message type: " + type.name());
}
diff --git a/samza-core/src/main/java/org/apache/samza/storage/SideInputTask.java b/samza-core/src/main/java/org/apache/samza/storage/SideInputTask.java
index 6274c15..981e847 100644
--- a/samza-core/src/main/java/org/apache/samza/storage/SideInputTask.java
+++ b/samza-core/src/main/java/org/apache/samza/storage/SideInputTask.java
@@ -97,6 +97,11 @@
}
@Override
+ public void drain(ReadableCoordinator coordinator) {
+ LOG.info("Task {} has drained", this.taskName);
+ }
+
+ @Override
public boolean isWindowableTask() {
return false;
}
diff --git a/samza-core/src/main/java/org/apache/samza/task/StreamOperatorTask.java b/samza-core/src/main/java/org/apache/samza/task/StreamOperatorTask.java
index 0079fab..0db960a 100644
--- a/samza-core/src/main/java/org/apache/samza/task/StreamOperatorTask.java
+++ b/samza-core/src/main/java/org/apache/samza/task/StreamOperatorTask.java
@@ -31,6 +31,7 @@
import org.apache.samza.system.EndOfStreamMessage;
import org.apache.samza.system.IncomingMessageEnvelope;
import org.apache.samza.system.MessageType;
+import org.apache.samza.system.DrainMessage;
import org.apache.samza.system.SystemStream;
import org.apache.samza.system.WatermarkMessage;
import org.apache.samza.util.Clock;
@@ -129,6 +130,12 @@
inputOpImpl.aggregateEndOfStream(eosMessage, ime.getSystemStreamPartition(), collector, coordinator);
break;
+ case DRAIN:
+ DrainMessage drainMessage = (DrainMessage) ime.getMessage();
+ processFuture =
+ inputOpImpl.aggregateDrainMessages(drainMessage, ime.getSystemStreamPartition(), collector, coordinator);
+ break;
+
case WATERMARK:
WatermarkMessage watermarkMessage = (WatermarkMessage) ime.getMessage();
processFuture = inputOpImpl.aggregateWatermark(watermarkMessage, ime.getSystemStreamPartition(), collector,
diff --git a/samza-core/src/main/scala/org/apache/samza/container/SamzaContainer.scala b/samza-core/src/main/scala/org/apache/samza/container/SamzaContainer.scala
index c157b87..f0aa20b 100644
--- a/samza-core/src/main/scala/org/apache/samza/container/SamzaContainer.scala
+++ b/samza-core/src/main/scala/org/apache/samza/container/SamzaContainer.scala
@@ -27,7 +27,6 @@
import java.util.concurrent._
import java.util.function.Consumer
import java.util.{Base64, Optional}
-import com.google.common.annotations.VisibleForTesting
import com.google.common.util.concurrent.ThreadFactoryBuilder
import org.apache.samza.SamzaException
import org.apache.samza.checkpoint.{CheckpointListener, OffsetManager, OffsetManagerMetrics}
@@ -426,6 +425,8 @@
val pollIntervalMs = taskConfig.getPollIntervalMs
+ val appConfig = new ApplicationConfig(config)
+
val consumerMultiplexer = new SystemConsumers(
chooser = chooser,
consumers = consumers,
@@ -435,7 +436,8 @@
dropDeserializationError = dropDeserializationError,
pollIntervalMs = pollIntervalMs,
clock = () => clock.nanoTime(),
- elasticityFactor = jobConfig.getElasticityFactor)
+ elasticityFactor = jobConfig.getElasticityFactor,
+ runId = appConfig.getRunId)
val producerMultiplexer = new SystemProducers(
producers = producers,
@@ -622,7 +624,7 @@
val maxThrottlingDelayMs = config.getLong("container.disk.quota.delay.max.ms", TimeUnit.SECONDS.toMillis(1))
- val runLoop = RunLoopFactory.createRunLoop(
+ val runLoop: Runnable = RunLoopFactory.createRunLoop(
taskInstances,
consumerMultiplexer,
taskThreadPool,
@@ -630,7 +632,8 @@
samzaContainerMetrics,
taskConfig,
clock,
- jobConfig.getElasticityFactor)
+ jobConfig.getElasticityFactor,
+ appConfig.getRunId)
val containerMemoryMb : Int = new ClusterManagerConfig(config).getContainerMemoryMb
@@ -747,6 +750,7 @@
def getStatus(): SamzaContainerStatus = status
def drain() {
+ // set the SystemConsumers multiplexer and RunLoop in drain mode
consumerMultiplexer.drain
}
diff --git a/samza-core/src/main/scala/org/apache/samza/container/TaskInstance.scala b/samza-core/src/main/scala/org/apache/samza/container/TaskInstance.scala
index d75d911..606f32a 100644
--- a/samza-core/src/main/scala/org/apache/samza/container/TaskInstance.scala
+++ b/samza-core/src/main/scala/org/apache/samza/container/TaskInstance.scala
@@ -72,6 +72,7 @@
val taskName: TaskName = taskModel.getTaskName
val isInitableTask = task.isInstanceOf[InitableTask]
+ val isDrainTask = task.isInstanceOf[DrainListenerTask]
val isEndOfStreamListenerTask = task.isInstanceOf[EndOfStreamListenerTask]
val isClosableTask = task.isInstanceOf[ClosableTask]
@@ -111,6 +112,12 @@
val streamConfig: StreamConfig = new StreamConfig(config)
override val intermediateStreams: java.util.Set[String] = JavaConverters.setAsJavaSetConverter(streamConfig.getStreamIds.filter(streamConfig.getIsIntermediateStream)).asJava
+ val intermediateSSPs: Set[SystemStreamPartition] = systemStreamPartitions.filter(ssp => {
+ val systemStream = ssp.getSystemStream
+ val streamId = streamConfig.systemStreamToStreamId(systemStream)
+ intermediateStreams.contains(streamId)
+ }).toSet
+
val streamsToDeleteCommittedMessages: Set[String] = streamConfig.getStreamIds.filter(streamConfig.getDeleteCommittedMessages).map(streamConfig.getPhysicalName).toSet
val checkpointWriteVersions = new TaskConfig(config).getCheckpointWriteVersions
@@ -189,20 +196,27 @@
}
/**
- * Computes the starting offset for the partitions assigned to the task and registers them with the underlying {@see SystemConsumers}.
- *
- * Starting offset for a partition of the task is computed in the following manner:
- *
- * 1. If a startpoint exists for a task, system stream partition and it resolves to a offset, then the resolved offset is used as the starting offset.
- * 2. Else, the checkpointed offset for the system stream partition is used as the starting offset.
+ * This method registers the following with the underlying {@see SystemConsumers}:
+ * a) starting offsets for all SSPs assigned to the task
+ * b) intermediate SSPs assigned to the task
+ *
+ * Starting offset for a partition of the task is computed in the following manner:
+ *
+ * 1. If a startpoint exists for a task, system stream partition and it resolves to a offset, then the resolved offset is used as the starting offset.
+ * 2. Else, the checkpointed offset for the system stream partition is used as the starting offset.
*/
def registerConsumers() {
debug("Registering consumers for taskName: %s" format taskName)
+
systemStreamPartitions.foreach(systemStreamPartition => {
val startingOffset: String = getStartingOffset(systemStreamPartition)
consumerMultiplexer.register(systemStreamPartition, startingOffset)
metrics.addOffsetGauge(systemStreamPartition, () => offsetManager.getLastProcessedOffset(taskName, systemStreamPartition).orNull)
})
+
+ intermediateSSPs.foreach(ssp => {
+ consumerMultiplexer.registerIntermediateSSP(ssp)
+ })
}
def process(envelope: IncomingMessageEnvelope, coordinator: ReadableCoordinator,
@@ -241,6 +255,16 @@
}
}
+ def drain(coordinator: ReadableCoordinator): Unit = {
+ task match {
+ case _: DrainListenerTask =>
+ exceptionHandler.maybeHandle {
+ task.asInstanceOf[DrainListenerTask].onDrain(collector, coordinator)
+ }
+ case _ =>
+ }
+ }
+
def window(coordinator: ReadableCoordinator) {
if (isWindowableTask) {
trace("Windowing for taskName: %s" format taskName)
@@ -613,14 +637,19 @@
// if elasticityFactor > 1, find the SSP with keyBucket
var incomingMessageSsp = envelope.getSystemStreamPartition(elasticityFactor)
- // if envelope is end of stream or watermark, it needs to be routed to all tasks consuming the ssp irresp of keyBucket
+ // if envelope is end of stream or watermark or drain,
+ // it needs to be routed to all tasks consuming the ssp irresp of keyBucket
val messageType = MessageType.of(envelope.getMessage)
- if (envelope.isEndOfStream() || MessageType.END_OF_STREAM == messageType || MessageType.WATERMARK == messageType) {
+ if (envelope.isEndOfStream()
+ || envelope.isDrain()
+ || messageType == MessageType.END_OF_STREAM
+ || messageType == MessageType.WATERMARK) {
+
incomingMessageSsp = systemStreamPartitions
.filter(ssp => ssp.getSystemStream.equals(incomingMessageSsp.getSystemStream)
&& ssp.getPartition.equals(incomingMessageSsp.getPartition))
.toIterator.next()
- debug("for watermark or end of stream envelope, found incoming ssp as {}" format incomingMessageSsp)
+ debug("for watermark or end-of-stream or drain envelope, found incoming ssp as {}".format(incomingMessageSsp))
}
incomingMessageSsp
}
diff --git a/samza-core/src/main/scala/org/apache/samza/job/local/ThreadJobFactory.scala b/samza-core/src/main/scala/org/apache/samza/job/local/ThreadJobFactory.scala
index f1c476a..ac1a558 100644
--- a/samza-core/src/main/scala/org/apache/samza/job/local/ThreadJobFactory.scala
+++ b/samza-core/src/main/scala/org/apache/samza/job/local/ThreadJobFactory.scala
@@ -89,7 +89,7 @@
}
var drainMonitor: DrainMonitor = null
if (jobConfig.getDrainMonitorEnabled()) {
- drainMonitor = new DrainMonitor(coordinatorStreamStore, config)
+ drainMonitor = new DrainMonitor(coordinatorStreamStore, config, jobConfig.getDrainMonitorPollIntervalMillis)
}
val containerId = "0"
diff --git a/samza-core/src/main/scala/org/apache/samza/storage/ContainerStorageManager.java b/samza-core/src/main/scala/org/apache/samza/storage/ContainerStorageManager.java
index c2e02a4..40b2999 100644
--- a/samza-core/src/main/scala/org/apache/samza/storage/ContainerStorageManager.java
+++ b/samza-core/src/main/scala/org/apache/samza/storage/ContainerStorageManager.java
@@ -45,6 +45,7 @@
import org.apache.samza.checkpoint.Checkpoint;
import org.apache.samza.checkpoint.CheckpointManager;
import org.apache.samza.checkpoint.CheckpointV2;
+import org.apache.samza.config.ApplicationConfig;
import org.apache.samza.config.Config;
import org.apache.samza.config.JobConfig;
import org.apache.samza.config.StorageConfig;
@@ -120,6 +121,8 @@
private static final int RESTORE_THREAD_POOL_SHUTDOWN_TIMEOUT_SECONDS = 60;
+ private static final int DEFAULT_SIDE_INPUT_ELASTICITY_FACTOR = 1;
+
/** Maps containing relevant per-task objects */
private final Map<TaskName, TaskInstanceMetrics> taskInstanceMetrics;
private final Map<TaskName, TaskInstanceCollector> taskInstanceCollectors;
@@ -297,11 +300,13 @@
MessageChooser chooser = DefaultChooser.apply(inputStreamMetadata, new RoundRobinChooserFactory(), config,
sideInputSystemConsumersMetrics.registry(), systemAdmins);
+ ApplicationConfig applicationConfig = new ApplicationConfig(config);
+
sideInputSystemConsumers =
new SystemConsumers(chooser, ScalaJavaUtil.toScalaMap(sideInputConsumers), systemAdmins, serdeManager,
sideInputSystemConsumersMetrics, SystemConsumers.DEFAULT_NO_NEW_MESSAGES_TIMEOUT(), SystemConsumers.DEFAULT_DROP_SERIALIZATION_ERROR(),
TaskConfig.DEFAULT_POLL_INTERVAL_MS, ScalaJavaUtil.toScalaFunction(() -> System.nanoTime()),
- JobConfig.DEFAULT_JOB_ELASTICITY_FACTOR);
+ JobConfig.DEFAULT_JOB_ELASTICITY_FACTOR, applicationConfig.getRunId());
}
}
@@ -922,6 +927,8 @@
new SamzaContainerMetrics(SIDEINPUTS_METRICS_PREFIX + this.samzaContainerMetrics.source(),
this.samzaContainerMetrics.registry(), SIDEINPUTS_METRICS_PREFIX);
+ ApplicationConfig applicationConfig = new ApplicationConfig(config);
+
this.sideInputRunLoop = new RunLoop(sideInputTasks,
null, // all operations are executed in the main runloop thread
this.sideInputSystemConsumers,
@@ -934,7 +941,9 @@
taskConfig.getMaxIdleMs(),
sideInputContainerMetrics,
System::nanoTime,
- false); // commit must be synchronous to ensure integrity of state flush
+ false,
+ DEFAULT_SIDE_INPUT_ELASTICITY_FACTOR,
+ applicationConfig.getRunId()); // commit must be synchronous to ensure integrity of state flush
try {
sideInputsExecutor.submit(() -> {
diff --git a/samza-core/src/main/scala/org/apache/samza/system/SystemConsumers.scala b/samza-core/src/main/scala/org/apache/samza/system/SystemConsumers.scala
index bcecebb..d1b639c 100644
--- a/samza-core/src/main/scala/org/apache/samza/system/SystemConsumers.scala
+++ b/samza-core/src/main/scala/org/apache/samza/system/SystemConsumers.scala
@@ -28,9 +28,8 @@
import java.util.HashSet
import java.util.Queue
import java.util.Set
-import java.util.function.Predicate
+import java.util.function.{Consumer}
import java.util.stream.Collectors
-
import scala.collection.JavaConverters._
import org.apache.samza.serializers.SerdeManager
import org.apache.samza.util.{Logging, TimerUtil}
@@ -38,6 +37,7 @@
import org.apache.samza.SamzaException
import org.apache.samza.config.TaskConfig
+
object SystemConsumers {
val DEFAULT_NO_NEW_MESSAGES_TIMEOUT = 10
val DEFAULT_DROP_SERIALIZATION_ERROR = false
@@ -116,7 +116,12 @@
*/
val clock: () => Long = () => System.nanoTime(),
- val elasticityFactor: Int = 1) extends Logging with TimerUtil {
+ val elasticityFactor: Int = 1,
+
+ /**
+ * Identifier of the current deployment.
+ * */
+ val runId: String = null) extends Logging with TimerUtil {
/**
* Mapping from the {@see SystemStreamPartition} to the registered offsets.
@@ -130,6 +135,10 @@
*/
private val sspKeyBucketsRegistered = new HashSet[SystemStreamPartition] ()
+ private val intermediateSSPs = new HashSet[SystemStreamPartition]()
+
+ private val intermediateSystems = new HashSet[String]()
+
/**
* A buffer of incoming messages grouped by SystemStreamPartition. These
* messages are handed out to the MessageChooser as it needs them.
@@ -155,10 +164,8 @@
*/
private var started = false
- /**
- * Denotes if the SystemConsumers is in drain mode.
- * */
- private var draining = false
+ @volatile
+ private var isDraining = false
/**
* Default timeout to noNewMessagesTimeout. Every time SystemConsumers
@@ -193,8 +200,6 @@
// but the actual systemConsumer which consumes from the input does not know about KeyBucket.
// hence, use an SSP without KeyBucket
consumer.register(removeKeyBucket(systemStreamPartition), offset)
- chooser.register(removeKeyBucket(systemStreamPartition), offset)
- debug("consumer.register and chooser.register for ssp: %s with offset %s" format (systemStreamPartition, offset))
}
debug("Starting consumers.")
@@ -214,15 +219,12 @@
chooser.start
+
started = true
refresh
}
- def drain: Unit = {
- draining = true
- }
-
def stop {
if (started) {
debug("Stopping consumers.")
@@ -237,6 +239,14 @@
}
}
+ def drain(): Unit = {
+ if (!isDraining) {
+ isDraining = true;
+ info("SystemConsumers is set to drain mode.")
+ consumers.values.foreach(_.stop)
+ writeDrainControlMessageToSspQueue()
+ }
+ }
def register(ssp: SystemStreamPartition, offset: String) {
// If elasticity is enabled then the RunLoop gives SSP with keybucket
@@ -255,6 +265,8 @@
metrics.registerSystemStreamPartition(systemStreamPartition)
unprocessedMessagesBySSP.put(systemStreamPartition, new ArrayDeque[IncomingMessageEnvelope]())
+ chooser.register(systemStreamPartition, offset)
+
try {
val consumer = consumers(systemStreamPartition.getSystem)
val existingOffset = sspToRegisteredOffsets.get(systemStreamPartition)
@@ -268,12 +280,17 @@
}
}
+ def registerIntermediateSSP(ssp: SystemStreamPartition): Unit = {
+ debug("Registering intermediate stream: %s" format ssp)
+ intermediateSSPs.add(ssp)
+ intermediateSystems.add(ssp.getSystem)
+ }
def isEndOfStream(systemStreamPartition: SystemStreamPartition) = {
endOfStreamSSPs.contains(removeKeyBucket(systemStreamPartition))
}
- def choose (updateChooser: Boolean = true): IncomingMessageEnvelope = {
+ def choose(updateChooser: Boolean = true): IncomingMessageEnvelope = {
val envelopeFromChooser = chooser.choose
updateTimer(metrics.deserializationNs) {
@@ -398,16 +415,54 @@
}
private def refresh {
- if (draining) {
- trace("Skipping refresh of chooser as the multiplexer is in drain mode.")
- return
- }
- trace("Refreshing chooser with new messages.")
-
// Update last poll time so we don't poll too frequently.
lastPollNs = clock()
- // Poll every system for new messages.
- consumers.keys.map(poll(_))
+
+ if (isDraining) {
+ trace("Refreshing chooser with new messages from intermediate systems.")
+
+ // scala 2.11 doesn't allow using syntactical sugar: intermediateSystems.foreach(poll(_)) over java collections
+ intermediateSystems.forEach(new Consumer[String] {
+ override def accept(system: String): Unit = poll(system)
+ })
+ } else {
+ trace("Refreshing chooser with new messages.")
+ consumers.keys.foreach(poll(_))
+ }
+ }
+
+ private def writeDrainControlMessageToSspQueue() {
+ val sspsToDrain = new HashSet(sspKeyBucketsRegistered)
+
+ // only write Drain ControlMessages to source SSPs
+ // sspsToDrain = allSSPs - intermediateSSPs - eosSSPs
+ sspsToDrain.removeAll(intermediateSSPs)
+ sspsToDrain.removeAll(endOfStreamSSPs)
+
+ sspsToDrain.forEach(new Consumer[SystemStreamPartition] {
+ override def accept(ssp: SystemStreamPartition): Unit = {
+ val envelopes: Queue[IncomingMessageEnvelope] =
+ if (unprocessedMessagesBySSP.containsKey(ssp)) {
+ unprocessedMessagesBySSP.get(ssp)
+ } else {
+ new util.ArrayDeque[IncomingMessageEnvelope]()
+ }
+
+ // Add watermark ControlMessage only if there are intermediate SSPs as low-level API task doesn't process
+ // WatermarkMessages
+ if (!intermediateSSPs.isEmpty) {
+ envelopes.add(IncomingMessageEnvelope.buildWatermarkEnvelope(ssp, Long.MaxValue))
+ totalUnprocessedMessages += 1
+ }
+ // Add Drain ControlMessage
+ envelopes.add(IncomingMessageEnvelope.buildDrainMessage(ssp, runId))
+ totalUnprocessedMessages += 1
+ unprocessedMessagesBySSP.put(ssp, envelopes)
+
+ // update the chooser with the messages
+ tryUpdate(ssp)
+ }
+ })
}
/**
diff --git a/samza-core/src/test/java/org/apache/samza/container/TestRunLoop.java b/samza-core/src/test/java/org/apache/samza/container/TestRunLoop.java
index 7bd5f9a..e584974 100644
--- a/samza-core/src/test/java/org/apache/samza/container/TestRunLoop.java
+++ b/samza-core/src/test/java/org/apache/samza/container/TestRunLoop.java
@@ -20,8 +20,9 @@
package org.apache.samza.container;
import com.google.common.collect.ImmutableMap;
-import java.util.Collections;
+import java.util.Arrays;
import java.util.HashMap;
+import java.util.HashSet;
import java.util.Map;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
@@ -49,6 +50,7 @@
public class TestRunLoop {
// Immutable objects shared by all test methods.
+ private final String runId = "foo";
private final ExecutorService executor = null;
private final SamzaContainerMetrics containerMetrics = new SamzaContainerMetrics("container", new MetricsRegistryMap(), "");
private final long windowMs = -1;
@@ -60,13 +62,23 @@
private final Partition p1 = new Partition(1);
private final TaskName taskName0 = new TaskName(p0.toString());
private final TaskName taskName1 = new TaskName(p1.toString());
- private final SystemStreamPartition ssp0 = new SystemStreamPartition("testSystem", "testStream", p0);
- private final SystemStreamPartition ssp1 = new SystemStreamPartition("testSystem", "testStream", p1);
- private final IncomingMessageEnvelope envelope00 = new IncomingMessageEnvelope(ssp0, "0", "key0", "value0");
- private final IncomingMessageEnvelope envelope11 = new IncomingMessageEnvelope(ssp1, "1", "key1", "value1");
- private final IncomingMessageEnvelope envelope01 = new IncomingMessageEnvelope(ssp0, "1", "key0", "value0");
- private final IncomingMessageEnvelope ssp0EndOfStream = IncomingMessageEnvelope.buildEndOfStreamEnvelope(ssp0);
- private final IncomingMessageEnvelope ssp1EndOfStream = IncomingMessageEnvelope.buildEndOfStreamEnvelope(ssp1);
+ private final SystemStreamPartition sspA0 = new SystemStreamPartition("testSystem", "testStreamA", p0);
+ private final SystemStreamPartition sspA1 = new SystemStreamPartition("testSystem", "testStreamA", p1);
+ private final SystemStreamPartition sspB0 = new SystemStreamPartition("testSystem", "testStreamB", p0);
+ private final SystemStreamPartition sspB1 = new SystemStreamPartition("testSystem", "testStreamB", p1);
+ private final IncomingMessageEnvelope envelopeA00 = new IncomingMessageEnvelope(sspA0, "0", "key0", "value0");
+ private final IncomingMessageEnvelope envelopeA11 = new IncomingMessageEnvelope(sspA1, "1", "key1", "value1");
+ private final IncomingMessageEnvelope envelopeA01 = new IncomingMessageEnvelope(sspA0, "1", "key0", "value0");
+ private final IncomingMessageEnvelope envelopeB00 = new IncomingMessageEnvelope(sspB0, "0", "key0", "value0");
+ private final IncomingMessageEnvelope envelopeB11 = new IncomingMessageEnvelope(sspB1, "1", "key1", "value1");
+ private final IncomingMessageEnvelope sspA0EndOfStream = IncomingMessageEnvelope.buildEndOfStreamEnvelope(sspA0);
+ private final IncomingMessageEnvelope sspA1EndOfStream = IncomingMessageEnvelope.buildEndOfStreamEnvelope(sspA1);
+ private final IncomingMessageEnvelope sspB0EndOfStream = IncomingMessageEnvelope.buildEndOfStreamEnvelope(sspB0);
+ private final IncomingMessageEnvelope sspB1EndOfStream = IncomingMessageEnvelope.buildEndOfStreamEnvelope(sspB1);
+ private final IncomingMessageEnvelope sspA0Drain = IncomingMessageEnvelope.buildDrainMessage(sspA0, runId);
+ private final IncomingMessageEnvelope sspA1Drain = IncomingMessageEnvelope.buildDrainMessage(sspA1, runId);
+ private final IncomingMessageEnvelope sspB0Drain = IncomingMessageEnvelope.buildDrainMessage(sspB0, runId);
+ private final IncomingMessageEnvelope sspB1Drain = IncomingMessageEnvelope.buildDrainMessage(sspB1, runId);
@Rule
public Timeout maxTestDurationInSeconds = Timeout.seconds(120);
@@ -75,8 +87,8 @@
public void testProcessMultipleTasks() {
SystemConsumers consumerMultiplexer = mock(SystemConsumers.class);
- RunLoopTask task0 = getMockRunLoopTask(taskName0, ssp0);
- RunLoopTask task1 = getMockRunLoopTask(taskName1, ssp1);
+ RunLoopTask task0 = getMockRunLoopTask(taskName0, sspA0);
+ RunLoopTask task1 = getMockRunLoopTask(taskName1, sspA1);
Map<TaskName, RunLoopTask> tasks = new HashMap<>();
tasks.put(taskName0, task0);
@@ -84,12 +96,13 @@
int maxMessagesInFlight = 1;
RunLoop runLoop = new RunLoop(tasks, executor, consumerMultiplexer, maxMessagesInFlight, windowMs, commitMs,
- callbackTimeoutMs, maxThrottlingDelayMs, maxIdleMs, containerMetrics, () -> 0L, false);
- when(consumerMultiplexer.choose(false)).thenReturn(envelope00).thenReturn(envelope11).thenReturn(ssp0EndOfStream).thenReturn(ssp1EndOfStream).thenReturn(null);
+ callbackTimeoutMs, maxThrottlingDelayMs, maxIdleMs, containerMetrics, () -> 0L, false, 1, "foo");
+ when(consumerMultiplexer.choose(false)).thenReturn(envelopeA00).thenReturn(envelopeA11).thenReturn(sspA0EndOfStream).thenReturn(
+ sspA1EndOfStream).thenReturn(null);
runLoop.run();
- verify(task0).process(eq(envelope00), any(), any());
- verify(task1).process(eq(envelope11), any(), any());
+ verify(task0).process(eq(envelopeA00), any(), any());
+ verify(task1).process(eq(envelopeA11), any(), any());
assertEquals(4L, containerMetrics.envelopes().getCount());
}
@@ -97,9 +110,9 @@
@Test
public void testProcessInOrder() {
SystemConsumers consumerMultiplexer = mock(SystemConsumers.class);
- when(consumerMultiplexer.choose(false)).thenReturn(envelope00).thenReturn(envelope01).thenReturn(ssp0EndOfStream).thenReturn(null);
+ when(consumerMultiplexer.choose(false)).thenReturn(envelopeA00).thenReturn(envelopeA01).thenReturn(sspA0EndOfStream).thenReturn(null);
- RunLoopTask task0 = getMockRunLoopTask(taskName0, ssp0);
+ RunLoopTask task0 = getMockRunLoopTask(taskName0, sspA0);
Map<TaskName, RunLoopTask> tasks = ImmutableMap.of(taskName0, task0);
int maxMessagesInFlight = 1;
@@ -108,8 +121,8 @@
runLoop.run();
InOrder inOrder = inOrder(task0);
- inOrder.verify(task0).process(eq(envelope00), any(), any());
- inOrder.verify(task0).process(eq(envelope01), any(), any());
+ inOrder.verify(task0).process(eq(envelopeA00), any(), any());
+ inOrder.verify(task0).process(eq(envelopeA01), any(), any());
}
@Test
@@ -119,7 +132,7 @@
SystemConsumers consumerMultiplexer = mock(SystemConsumers.class);
OffsetManager offsetManager = mock(OffsetManager.class);
- RunLoopTask task0 = getMockRunLoopTask(taskName0, ssp0);
+ RunLoopTask task0 = getMockRunLoopTask(taskName0, sspA0);
when(task0.offsetManager()).thenReturn(offsetManager);
CountDownLatch firstMessageBarrier = new CountDownLatch(1);
doAnswer(invocation -> {
@@ -134,7 +147,7 @@
return null;
});
return null;
- }).when(task0).process(eq(envelope00), any(), any());
+ }).when(task0).process(eq(envelopeA00), any(), any());
doAnswer(invocation -> {
assertEquals(1, task0.metrics().messagesInFlight().getValue());
@@ -145,21 +158,21 @@
callback.complete();
firstMessageBarrier.countDown();
return null;
- }).when(task0).process(eq(envelope01), any(), any());
+ }).when(task0).process(eq(envelopeA01), any(), any());
Map<TaskName, RunLoopTask> tasks = new HashMap<>();
tasks.put(taskName0, task0);
RunLoop runLoop = new RunLoop(tasks, executor, consumerMultiplexer, maxMessagesInFlight, windowMs, commitMs,
callbackTimeoutMs, maxThrottlingDelayMs, maxIdleMs, containerMetrics, () -> 0L, false);
- when(consumerMultiplexer.choose(false)).thenReturn(envelope00).thenReturn(envelope01).thenReturn(null);
+ when(consumerMultiplexer.choose(false)).thenReturn(envelopeA00).thenReturn(envelopeA01).thenReturn(null);
runLoop.run();
InOrder inOrder = inOrder(task0);
- inOrder.verify(task0).process(eq(envelope00), any(), any());
- inOrder.verify(task0).process(eq(envelope01), any(), any());
+ inOrder.verify(task0).process(eq(envelopeA00), any(), any());
+ inOrder.verify(task0).process(eq(envelopeA01), any(), any());
- verify(offsetManager).update(eq(taskName0), eq(ssp0), eq(envelope00.getOffset()));
+ verify(offsetManager).update(eq(taskName0), eq(sspA0), eq(envelopeA00.getOffset()));
assertEquals(2L, containerMetrics.processes().getCount());
}
@@ -168,9 +181,9 @@
public void testProcessElasticityEnabled() {
TaskName taskName0 = new TaskName(p0.toString() + " 0");
- SystemStreamPartition ssp = new SystemStreamPartition("testSystem", "testStream", p0);
- SystemStreamPartition ssp0 = new SystemStreamPartition("testSystem", "testStream", p0, 0);
- SystemStreamPartition ssp1 = new SystemStreamPartition("testSystem", "testStream", p0, 1);
+ SystemStreamPartition ssp = new SystemStreamPartition("testSystem", "testStreamA", p0);
+ SystemStreamPartition ssp0 = new SystemStreamPartition("testSystem", "testStreamA", p0, 0);
+ SystemStreamPartition ssp1 = new SystemStreamPartition("testSystem", "testStreamA", p0, 1);
// create two IME such that one of their ssp keybucket maps to ssp0 and the other one maps to ssp1
// task in the runloop should process only the first ime (aka the one whose ssp keybucket is ssp0)
@@ -197,12 +210,12 @@
}).when(task0).process(eq(envelope01), any(), any());
SystemConsumers consumerMultiplexer = mock(SystemConsumers.class);
- when(consumerMultiplexer.choose(false)).thenReturn(envelope00).thenReturn(envelope01).thenReturn(ssp0EndOfStream).thenReturn(null);
+ when(consumerMultiplexer.choose(false)).thenReturn(envelope00).thenReturn(envelope01).thenReturn(sspA0EndOfStream).thenReturn(null);
Map<TaskName, RunLoopTask> tasks = ImmutableMap.of(taskName0, task0);
int maxMessagesInFlight = 1;
RunLoop runLoop = new RunLoop(tasks, executor, consumerMultiplexer, maxMessagesInFlight, windowMs, commitMs,
- callbackTimeoutMs, maxThrottlingDelayMs, 0, containerMetrics, () -> 0L, false, 2);
+ callbackTimeoutMs, maxThrottlingDelayMs, 0, containerMetrics, () -> 0L, false, 2, null);
runLoop.run();
verify(task0).process(eq(envelope00), any(), any());
@@ -216,9 +229,9 @@
TaskName taskName0 = new TaskName(p0.toString() + " 0");
TaskName taskName1 = new TaskName(p0.toString() + " 1");
- SystemStreamPartition ssp = new SystemStreamPartition("testSystem", "testStream", p0);
- SystemStreamPartition ssp0 = new SystemStreamPartition("testSystem", "testStream", p0, 0);
- SystemStreamPartition ssp1 = new SystemStreamPartition("testSystem", "testStream", p0, 1);
+ SystemStreamPartition ssp = new SystemStreamPartition("testSystem", "testStreamA", p0);
+ SystemStreamPartition ssp0 = new SystemStreamPartition("testSystem", "testStreamA", p0, 0);
+ SystemStreamPartition ssp1 = new SystemStreamPartition("testSystem", "testStreamA", p0, 1);
// create EOS IME such that its ssp keybucket maps to ssp0 and not to ssp1
// task in the runloop should give this ime to both it tasks
@@ -237,7 +250,7 @@
Map<TaskName, RunLoopTask> tasks = ImmutableMap.of(taskName0, task0, taskName1, task1);
int maxMessagesInFlight = 1;
RunLoop runLoop = new RunLoop(tasks, executor, consumerMultiplexer, maxMessagesInFlight, windowMs, commitMs,
- callbackTimeoutMs, maxThrottlingDelayMs, 0, containerMetrics, () -> 0L, false, 2);
+ callbackTimeoutMs, maxThrottlingDelayMs, 0, containerMetrics, () -> 0L, false, 2, null);
runLoop.run();
verify(task0).endOfStream(any());
@@ -245,12 +258,104 @@
}
@Test
+ public void testDrainWithElasticityEnabled() {
+ TaskName taskName0 = new TaskName(p0.toString() + " 0");
+ TaskName taskName1 = new TaskName(p0.toString() + " 1");
+ SystemStreamPartition ssp = new SystemStreamPartition("testSystem", "testStreamA", p0);
+ SystemStreamPartition ssp0 = new SystemStreamPartition("testSystem", "testStreamA", p0, 0);
+ SystemStreamPartition ssp1 = new SystemStreamPartition("testSystem", "testStreamA", p0, 1);
+
+ // create EOS IME such that its ssp keybucket maps to ssp0 and not to ssp1
+ // task in the runloop should give this ime to both it tasks
+ IncomingMessageEnvelope envelopeDrain = spy(IncomingMessageEnvelope.buildDrainMessage(ssp, runId));
+ when(envelopeDrain.getSystemStreamPartition(2)).thenReturn(ssp0);
+
+ // two task in the run loop that processes ssp0 -> 0th keybucket of ssp and ssp1 -> 1st keybucket of ssp
+ // Drain ime should be given to both the tasks irrespective of the keybucket
+ RunLoopTask task0 = getMockRunLoopTask(taskName0, ssp0);
+ RunLoopTask task1 = getMockRunLoopTask(taskName1, ssp1);
+
+ SystemConsumers consumerMultiplexer = mock(SystemConsumers.class);
+ when(consumerMultiplexer.choose(false)).thenReturn(envelopeDrain).thenReturn(null);
+
+ Map<TaskName, RunLoopTask> tasks = ImmutableMap.of(taskName0, task0, taskName1, task1);
+ int maxMessagesInFlight = 1;
+ RunLoop runLoop = new RunLoop(tasks, executor, consumerMultiplexer, maxMessagesInFlight, windowMs, commitMs,
+ callbackTimeoutMs, maxThrottlingDelayMs, 0, containerMetrics, () -> 0L, false, 2, runId);
+ runLoop.run();
+
+ verify(task0).drain(any());
+ verify(task1).drain(any());
+ }
+
+
+ @Test
+ public void testDrainForTasksWithSingleSSP() {
+ TaskName taskName0 = new TaskName(p0.toString() + " 0");
+ TaskName taskName1 = new TaskName(p1.toString() + " 1");
+
+ RunLoopTask task0 = getMockRunLoopTask(taskName0, sspA0);
+ RunLoopTask task1 = getMockRunLoopTask(taskName1, sspA1);
+
+ SystemConsumers consumerMultiplexer = mock(SystemConsumers.class);
+ // insert all envelopes followed by drain messages
+ when(consumerMultiplexer.choose(false))
+ .thenReturn(envelopeA00).thenReturn(envelopeA11)
+ .thenReturn(sspA0Drain).thenReturn(sspA1Drain);
+
+ Map<TaskName, RunLoopTask> tasks = ImmutableMap.of(taskName0, task0, taskName1, task1);
+ int maxMessagesInFlight = 1;
+ RunLoop runLoop = new RunLoop(tasks, executor, consumerMultiplexer, maxMessagesInFlight, windowMs, commitMs,
+ callbackTimeoutMs, maxThrottlingDelayMs, 0, containerMetrics, () -> 0L, false, 1, runId);
+ runLoop.run();
+
+ // check if process was called once for each task
+ verify(task0, times(1)).process(any(), any(), any());
+ verify(task1, times(1)).process(any(), any(), any());
+ // check if drain was called once for each task followed by commit
+ verify(task0, times(1)).drain(any());
+ verify(task1, times(1)).drain(any());
+ verify(task0, times(1)).commit();
+ verify(task1, times(1)).commit();
+ }
+
+ @Test
+ public void testDrainForTasksWithMultipleSSP() {
+ TaskName taskName0 = new TaskName(p0.toString() + " 0");
+ TaskName taskName1 = new TaskName(p1.toString() + " 1");
+
+ RunLoopTask task0 = getMockRunLoopTask(taskName0, sspA0, sspB0);
+ RunLoopTask task1 = getMockRunLoopTask(taskName1, sspA1, sspB1);
+
+ SystemConsumers consumerMultiplexer = mock(SystemConsumers.class);
+ // insert all envelopes followed by drain messages
+ when(consumerMultiplexer.choose(false))
+ .thenReturn(envelopeA00).thenReturn(envelopeA11).thenReturn(envelopeB00).thenReturn(envelopeB11)
+ .thenReturn(sspA0Drain).thenReturn(sspA1Drain).thenReturn(sspB0Drain).thenReturn(sspB1Drain);
+
+ Map<TaskName, RunLoopTask> tasks = ImmutableMap.of(taskName0, task0, taskName1, task1);
+ int maxMessagesInFlight = 1;
+ RunLoop runLoop = new RunLoop(tasks, executor, consumerMultiplexer, maxMessagesInFlight, windowMs, commitMs,
+ callbackTimeoutMs, maxThrottlingDelayMs, 0, containerMetrics, () -> 0L, false, 1, runId);
+ runLoop.run();
+
+ // check if process was called twice for each task
+ verify(task0, times(2)).process(any(), any(), any());
+ verify(task1, times(2)).process(any(), any(), any());
+ // check if drain was called once for each task followed by commit
+ verify(task0, times(1)).drain(any());
+ verify(task1, times(1)).drain(any());
+ verify(task0, times(1)).commit();
+ verify(task1, times(1)).commit();
+ }
+
+ @Test
public void testWindow() {
SystemConsumers consumerMultiplexer = mock(SystemConsumers.class);
int maxMessagesInFlight = 1;
long windowMs = 1;
- RunLoopTask task = getMockRunLoopTask(taskName0, ssp0);
+ RunLoopTask task = getMockRunLoopTask(taskName0, sspA0);
when(task.isWindowableTask()).thenReturn(true);
final AtomicInteger windowCount = new AtomicInteger(0);
@@ -277,7 +382,7 @@
public void testCommitSingleTask() {
SystemConsumers consumerMultiplexer = mock(SystemConsumers.class);
- RunLoopTask task0 = getMockRunLoopTask(taskName0, ssp0);
+ RunLoopTask task0 = getMockRunLoopTask(taskName0, sspA0);
doAnswer(invocation -> {
ReadableCoordinator coordinator = invocation.getArgumentAt(1, ReadableCoordinator.class);
TaskCallbackFactory callbackFactory = invocation.getArgumentAt(2, TaskCallbackFactory.class);
@@ -288,9 +393,9 @@
callback.complete();
return null;
- }).when(task0).process(eq(envelope00), any(), any());
+ }).when(task0).process(eq(envelopeA00), any(), any());
- RunLoopTask task1 = getMockRunLoopTask(taskName1, ssp1);
+ RunLoopTask task1 = getMockRunLoopTask(taskName1, sspA1);
Map<TaskName, RunLoopTask> tasks = new HashMap<>();
tasks.put(this.taskName0, task0);
@@ -300,7 +405,7 @@
RunLoop runLoop = new RunLoop(tasks, executor, consumerMultiplexer, maxMessagesInFlight, windowMs, commitMs,
callbackTimeoutMs, maxThrottlingDelayMs, maxIdleMs, containerMetrics, () -> 0L, false);
//have a null message in between to make sure task0 finishes processing and invoke the commit
- when(consumerMultiplexer.choose(false)).thenReturn(envelope00).thenReturn(envelope11).thenReturn(null);
+ when(consumerMultiplexer.choose(false)).thenReturn(envelopeA00).thenReturn(envelopeA11).thenReturn(null);
runLoop.run();
@@ -315,7 +420,7 @@
public void testCommitAllTasks() {
SystemConsumers consumerMultiplexer = mock(SystemConsumers.class);
- RunLoopTask task0 = getMockRunLoopTask(taskName0, ssp0);
+ RunLoopTask task0 = getMockRunLoopTask(taskName0, sspA0);
doAnswer(invocation -> {
ReadableCoordinator coordinator = invocation.getArgumentAt(1, ReadableCoordinator.class);
TaskCallbackFactory callbackFactory = invocation.getArgumentAt(2, TaskCallbackFactory.class);
@@ -326,9 +431,9 @@
callback.complete();
return null;
- }).when(task0).process(eq(envelope00), any(), any());
+ }).when(task0).process(eq(envelopeA00), any(), any());
- RunLoopTask task1 = getMockRunLoopTask(taskName1, ssp1);
+ RunLoopTask task1 = getMockRunLoopTask(taskName1, sspA1);
Map<TaskName, RunLoopTask> tasks = new HashMap<>();
tasks.put(this.taskName0, task0);
@@ -338,7 +443,7 @@
RunLoop runLoop = new RunLoop(tasks, executor, consumerMultiplexer, maxMessagesInFlight, windowMs, commitMs,
callbackTimeoutMs, maxThrottlingDelayMs, maxIdleMs, containerMetrics, () -> 0L, false);
//have a null message in between to make sure task0 finishes processing and invoke the commit
- when(consumerMultiplexer.choose(false)).thenReturn(envelope00).thenReturn(envelope11).thenReturn(null);
+ when(consumerMultiplexer.choose(false)).thenReturn(envelopeA00).thenReturn(envelopeA11).thenReturn(null);
runLoop.run();
@@ -354,7 +459,7 @@
SystemConsumers consumerMultiplexer = mock(SystemConsumers.class);
int maxMessagesInFlight = 1;
- RunLoopTask task0 = getMockRunLoopTask(taskName0, ssp0);
+ RunLoopTask task0 = getMockRunLoopTask(taskName0, sspA0);
doAnswer(invocation -> {
ReadableCoordinator coordinator = invocation.getArgumentAt(1, ReadableCoordinator.class);
TaskCallbackFactory callbackFactory = invocation.getArgumentAt(2, TaskCallbackFactory.class);
@@ -363,9 +468,9 @@
coordinator.shutdown(TaskCoordinator.RequestScope.CURRENT_TASK);
callback.complete();
return null;
- }).when(task0).process(eq(envelope00), any(), any());
+ }).when(task0).process(eq(envelopeA00), any(), any());
- RunLoopTask task1 = getMockRunLoopTask(taskName1, ssp1);
+ RunLoopTask task1 = getMockRunLoopTask(taskName1, sspA1);
doAnswer(invocation -> {
ReadableCoordinator coordinator = invocation.getArgumentAt(1, ReadableCoordinator.class);
TaskCallbackFactory callbackFactory = invocation.getArgumentAt(2, TaskCallbackFactory.class);
@@ -374,7 +479,7 @@
coordinator.shutdown(TaskCoordinator.RequestScope.CURRENT_TASK);
callback.complete();
return null;
- }).when(task1).process(eq(envelope11), any(), any());
+ }).when(task1).process(eq(envelopeA11), any(), any());
Map<TaskName, RunLoopTask> tasks = new HashMap<>();
tasks.put(taskName0, task0);
@@ -383,7 +488,7 @@
RunLoop runLoop = new RunLoop(tasks, executor, consumerMultiplexer, maxMessagesInFlight, windowMs, commitMs,
callbackTimeoutMs, maxThrottlingDelayMs, maxIdleMs, containerMetrics, () -> 0L, false);
// consensus is reached after envelope1 is processed.
- when(consumerMultiplexer.choose(false)).thenReturn(envelope00).thenReturn(envelope11).thenReturn(null);
+ when(consumerMultiplexer.choose(false)).thenReturn(envelopeA00).thenReturn(envelopeA11).thenReturn(null);
runLoop.run();
verify(task0).process(any(), any(), any());
@@ -397,8 +502,8 @@
public void testEndOfStreamWithMultipleTasks() {
SystemConsumers consumerMultiplexer = mock(SystemConsumers.class);
- RunLoopTask task0 = getMockRunLoopTask(taskName0, ssp0);
- RunLoopTask task1 = getMockRunLoopTask(taskName1, ssp1);
+ RunLoopTask task0 = getMockRunLoopTask(taskName0, sspA0, sspB0);
+ RunLoopTask task1 = getMockRunLoopTask(taskName1, sspA1, sspB1);
Map<TaskName, RunLoopTask> tasks = new HashMap<>();
@@ -409,21 +514,27 @@
RunLoop runLoop = new RunLoop(tasks, executor, consumerMultiplexer, maxMessagesInFlight, windowMs, commitMs,
callbackTimeoutMs, maxThrottlingDelayMs, maxIdleMs, containerMetrics, () -> 0L, false);
when(consumerMultiplexer.choose(false))
- .thenReturn(envelope00)
- .thenReturn(envelope11)
- .thenReturn(ssp0EndOfStream)
- .thenReturn(ssp1EndOfStream)
+ .thenReturn(envelopeA00)
+ .thenReturn(envelopeA11)
+ .thenReturn(envelopeB00)
+ .thenReturn(envelopeB11)
+ .thenReturn(sspA0EndOfStream)
+ .thenReturn(sspB0EndOfStream)
+ .thenReturn(sspB1EndOfStream)
+ .thenReturn(sspA1EndOfStream)
.thenReturn(null);
runLoop.run();
- verify(task0).process(eq(envelope00), any(), any());
+ verify(task0).process(eq(envelopeA00), any(), any());
+ verify(task0).process(eq(envelopeB00), any(), any());
verify(task0).endOfStream(any());
- verify(task1).process(eq(envelope11), any(), any());
+ verify(task1).process(eq(envelopeA11), any(), any());
+ verify(task1).process(eq(envelopeB11), any(), any());
verify(task1).endOfStream(any());
- assertEquals(4L, containerMetrics.envelopes().getCount());
+ assertEquals(8L, containerMetrics.envelopes().getCount());
}
@Test
@@ -433,7 +544,7 @@
SystemConsumers consumerMultiplexer = mock(SystemConsumers.class);
OffsetManager offsetManager = mock(OffsetManager.class);
- RunLoopTask task0 = getMockRunLoopTask(taskName0, ssp0);
+ RunLoopTask task0 = getMockRunLoopTask(taskName0, sspA0);
when(task0.offsetManager()).thenReturn(offsetManager);
CountDownLatch firstMessageBarrier = new CountDownLatch(2);
doAnswer(invocation -> {
@@ -445,7 +556,7 @@
return null;
});
return null;
- }).when(task0).process(eq(envelope00), any(), any());
+ }).when(task0).process(eq(envelopeA00), any(), any());
doAnswer(invocation -> {
assertEquals(1, task0.metrics().messagesInFlight().getValue());
@@ -455,7 +566,7 @@
callback.complete();
firstMessageBarrier.countDown();
return null;
- }).when(task0).process(eq(envelope01), any(), any());
+ }).when(task0).process(eq(envelopeA01), any(), any());
doAnswer(invocation -> {
assertEquals(0, task0.metrics().messagesInFlight().getValue());
@@ -469,7 +580,7 @@
RunLoop runLoop = new RunLoop(tasks, executor, consumerMultiplexer, maxMessagesInFlight, windowMs, commitMs,
callbackTimeoutMs, maxThrottlingDelayMs, maxIdleMs, containerMetrics, () -> 0L, false);
- when(consumerMultiplexer.choose(false)).thenReturn(envelope00).thenReturn(envelope01).thenReturn(ssp0EndOfStream)
+ when(consumerMultiplexer.choose(false)).thenReturn(envelopeA00).thenReturn(envelopeA01).thenReturn(sspA0EndOfStream)
.thenAnswer(invocation -> {
// this ensures that the end of stream message has passed through run loop BEFORE the last remaining in flight message completes
firstMessageBarrier.countDown();
@@ -482,10 +593,67 @@
}
@Test
+ public void testDrainWaitsForInFlightMessages() {
+ int maxMessagesInFlight = 2;
+ ExecutorService taskExecutor = Executors.newFixedThreadPool(1);
+ SystemConsumers consumerMultiplexer = mock(SystemConsumers.class);
+ OffsetManager offsetManager = mock(OffsetManager.class);
+
+ RunLoopTask task0 = getMockRunLoopTask(taskName0, sspA0);
+ when(task0.offsetManager()).thenReturn(offsetManager);
+ CountDownLatch firstMessageBarrier = new CountDownLatch(2);
+ doAnswer(invocation -> {
+ TaskCallbackFactory callbackFactory = invocation.getArgumentAt(2, TaskCallbackFactory.class);
+ TaskCallback callback = callbackFactory.createCallback();
+ taskExecutor.submit(() -> {
+ firstMessageBarrier.await();
+ callback.complete();
+ return null;
+ });
+ return null;
+ }).when(task0).process(eq(envelopeA00), any(), any());
+
+ doAnswer(invocation -> {
+ assertEquals(1, task0.metrics().messagesInFlight().getValue());
+
+ TaskCallbackFactory callbackFactory = invocation.getArgumentAt(2, TaskCallbackFactory.class);
+ TaskCallback callback = callbackFactory.createCallback();
+ callback.complete();
+ firstMessageBarrier.countDown();
+ return null;
+ }).when(task0).process(eq(envelopeA01), any(), any());
+
+ doAnswer(invocation -> {
+ assertEquals(0, task0.metrics().messagesInFlight().getValue());
+ assertEquals(2, task0.metrics().asyncCallbackCompleted().getCount());
+
+ return null;
+ }).when(task0).drain(any());
+
+ Map<TaskName, RunLoopTask> tasks = new HashMap<>();
+ tasks.put(taskName0, task0);
+
+ RunLoop runLoop = new RunLoop(tasks, executor, consumerMultiplexer, maxMessagesInFlight, windowMs, commitMs,
+ callbackTimeoutMs, maxThrottlingDelayMs, maxIdleMs, containerMetrics, () -> 0L, false,
+ 1, runId);
+ when(consumerMultiplexer.choose(false)).thenReturn(envelopeA00).thenReturn(envelopeA01).thenReturn(sspA0Drain)
+ .thenAnswer(invocation -> {
+ // this ensures that the drain message has passed through run loop BEFORE the flight message
+ // completes
+ firstMessageBarrier.countDown();
+ return null;
+ });
+
+ runLoop.run();
+
+ verify(task0).drain(any());
+ }
+
+ @Test
public void testEndOfStreamCommitBehavior() {
SystemConsumers consumerMultiplexer = mock(SystemConsumers.class);
- RunLoopTask task0 = getMockRunLoopTask(taskName0, ssp0);
+ RunLoopTask task0 = getMockRunLoopTask(taskName0, sspA0);
doAnswer(invocation -> {
ReadableCoordinator coordinator = invocation.getArgumentAt(0, ReadableCoordinator.class);
@@ -500,7 +668,7 @@
int maxMessagesInFlight = 1;
RunLoop runLoop = new RunLoop(tasks, executor, consumerMultiplexer, maxMessagesInFlight, windowMs, commitMs,
callbackTimeoutMs, maxThrottlingDelayMs, maxIdleMs, containerMetrics, () -> 0L, false);
- when(consumerMultiplexer.choose(false)).thenReturn(envelope00).thenReturn(ssp0EndOfStream).thenReturn(null);
+ when(consumerMultiplexer.choose(false)).thenReturn(envelopeA00).thenReturn(sspA0EndOfStream).thenReturn(null);
runLoop.run();
@@ -511,13 +679,36 @@
}
@Test
+ public void testDrainCommitBehavior() {
+ SystemConsumers consumerMultiplexer = mock(SystemConsumers.class);
+
+ RunLoopTask task0 = getMockRunLoopTask(taskName0, sspA0);
+ Map<TaskName, RunLoopTask> tasks = new HashMap<>();
+
+ tasks.put(taskName0, task0);
+
+ int maxMessagesInFlight = 1;
+ RunLoop runLoop = new RunLoop(tasks, executor, consumerMultiplexer, maxMessagesInFlight, windowMs, commitMs,
+ callbackTimeoutMs, maxThrottlingDelayMs, maxIdleMs, containerMetrics, () -> 0L, false,
+ 1, runId);
+ when(consumerMultiplexer.choose(false)).thenReturn(envelopeA00).thenReturn(sspA0Drain).thenReturn(null);
+
+ runLoop.run();
+
+ InOrder inOrder = inOrder(task0);
+
+ inOrder.verify(task0).drain(any());
+ inOrder.verify(task0).commit();
+ }
+
+ @Test
public void testCommitWithMessageInFlightWhenAsyncCommitIsEnabled() {
int maxMessagesInFlight = 2;
ExecutorService taskExecutor = Executors.newFixedThreadPool(2);
SystemConsumers consumerMultiplexer = mock(SystemConsumers.class);
OffsetManager offsetManager = mock(OffsetManager.class);
- RunLoopTask task0 = getMockRunLoopTask(taskName0, ssp0);
+ RunLoopTask task0 = getMockRunLoopTask(taskName0, sspA0);
when(task0.offsetManager()).thenReturn(offsetManager);
CountDownLatch firstMessageBarrier = new CountDownLatch(1);
doAnswer(invocation -> {
@@ -532,7 +723,7 @@
return null;
});
return null;
- }).when(task0).process(eq(envelope00), any(), any());
+ }).when(task0).process(eq(envelopeA00), any(), any());
CountDownLatch secondMessageBarrier = new CountDownLatch(1);
doAnswer(invocation -> {
@@ -550,7 +741,7 @@
return null;
});
return null;
- }).when(task0).process(eq(envelope01), any(), any());
+ }).when(task0).process(eq(envelopeA01), any(), any());
doAnswer(invocation -> {
assertEquals(1, task0.metrics().asyncCallbackCompleted().getCount());
@@ -565,12 +756,12 @@
RunLoop runLoop = new RunLoop(tasks, executor, consumerMultiplexer, maxMessagesInFlight, windowMs, commitMs,
callbackTimeoutMs, maxThrottlingDelayMs, maxIdleMs, containerMetrics, () -> 0L, true);
- when(consumerMultiplexer.choose(false)).thenReturn(envelope00).thenReturn(envelope01).thenReturn(null);
+ when(consumerMultiplexer.choose(false)).thenReturn(envelopeA00).thenReturn(envelopeA01).thenReturn(null);
runLoop.run();
InOrder inOrder = inOrder(task0);
- inOrder.verify(task0).process(eq(envelope00), any(), any());
- inOrder.verify(task0).process(eq(envelope01), any(), any());
+ inOrder.verify(task0).process(eq(envelopeA00), any(), any());
+ inOrder.verify(task0).process(eq(envelopeA01), any(), any());
inOrder.verify(task0).commit();
}
@@ -578,12 +769,12 @@
public void testExceptionIsPropagated() {
SystemConsumers consumerMultiplexer = mock(SystemConsumers.class);
- RunLoopTask task0 = getMockRunLoopTask(taskName0, ssp0);
+ RunLoopTask task0 = getMockRunLoopTask(taskName0, sspA0);
doAnswer(invocation -> {
TaskCallbackFactory callbackFactory = invocation.getArgumentAt(2, TaskCallbackFactory.class);
callbackFactory.createCallback().failure(new Exception("Intentional failure"));
return null;
- }).when(task0).process(eq(envelope00), any(), any());
+ }).when(task0).process(eq(envelopeA00), any(), any());
Map<TaskName, RunLoopTask> tasks = ImmutableMap.of(taskName0, task0);
@@ -592,16 +783,16 @@
callbackTimeoutMs, maxThrottlingDelayMs, maxIdleMs, containerMetrics, () -> 0L, false);
when(consumerMultiplexer.choose(false))
- .thenReturn(envelope00)
- .thenReturn(ssp0EndOfStream)
+ .thenReturn(envelopeA00)
+ .thenReturn(sspA0EndOfStream)
.thenReturn(null);
runLoop.run();
}
- private RunLoopTask getMockRunLoopTask(TaskName taskName, SystemStreamPartition ssp0) {
+ private RunLoopTask getMockRunLoopTask(TaskName taskName, SystemStreamPartition ... ssps) {
RunLoopTask task0 = mock(RunLoopTask.class);
- when(task0.systemStreamPartitions()).thenReturn(Collections.singleton(ssp0));
+ when(task0.systemStreamPartitions()).thenReturn(new HashSet<>(Arrays.asList(ssps)));
when(task0.metrics()).thenReturn(new TaskInstanceMetrics("test", new MetricsRegistryMap(), ""));
when(task0.taskName()).thenReturn(taskName);
return task0;
diff --git a/samza-core/src/test/java/org/apache/samza/drain/DrainMonitorTests.java b/samza-core/src/test/java/org/apache/samza/drain/DrainMonitorTests.java
index e7666aa..80061b3 100644
--- a/samza-core/src/test/java/org/apache/samza/drain/DrainMonitorTests.java
+++ b/samza-core/src/test/java/org/apache/samza/drain/DrainMonitorTests.java
@@ -41,13 +41,13 @@
* Tests for {@link DrainMonitor}
* */
public class DrainMonitorTests {
- private static final String TEST_DEPLOYMENT_ID = "foo";
+ private static final String TEST_RUN_ID = "foo";
private static final Config
CONFIG = new MapConfig(ImmutableMap.of(
"job.name", "test-job",
"job.coordinator.system", "test-kafka",
- ApplicationConfig.APP_RUN_ID, TEST_DEPLOYMENT_ID));
+ ApplicationConfig.APP_RUN_ID, TEST_RUN_ID));
private CoordinatorStreamStore coordinatorStreamStore;
@@ -114,7 +114,7 @@
final AtomicInteger numCallbacks = new AtomicInteger(0);
final CountDownLatch latch = new CountDownLatch(1);
// write drain before monitor start
- DrainUtils.writeDrainNotification(coordinatorStreamStore, TEST_DEPLOYMENT_ID);
+ DrainUtils.writeDrainNotification(coordinatorStreamStore, TEST_RUN_ID);
DrainMonitor drainMonitor = new DrainMonitor(coordinatorStreamStore, CONFIG);
drainMonitor.registerDrainCallback(() -> {
numCallbacks.incrementAndGet();
@@ -141,7 +141,7 @@
latch.countDown();
});
drainMonitor.start();
- DrainUtils.writeDrainNotification(coordinatorStreamStore, TEST_DEPLOYMENT_ID);
+ DrainUtils.writeDrainNotification(coordinatorStreamStore, TEST_RUN_ID);
if (!latch.await(2, TimeUnit.SECONDS)) {
Assert.fail("Timed out waiting for drain callback to complete");
}
@@ -150,8 +150,8 @@
}
@Test
- public void testCallbackNotCalledDueToMismatchedDeploymentId() throws InterruptedException {
- // The test fails due to timeout as the published DrainNotification's deploymentId doesn't match deploymentId
+ public void testCallbackNotCalledDueToMismatchedRunId() throws InterruptedException {
+ // The test fails due to timeout as the published DrainNotification's runId doesn't match runId
// in the config
exceptionRule.expect(AssertionError.class);
exceptionRule.expectMessage("Timed out waiting for drain callback to complete.");
@@ -166,8 +166,8 @@
});
drainMonitor.start();
- final String mismatchedDeploymentId = "bar";
- DrainUtils.writeDrainNotification(coordinatorStreamStore, mismatchedDeploymentId);
+ final String mismatchedRunId = "bar";
+ DrainUtils.writeDrainNotification(coordinatorStreamStore, mismatchedRunId);
if (!latch.await(2, TimeUnit.SECONDS)) {
Assert.fail("Timed out waiting for drain callback to complete.");
}
@@ -184,16 +184,16 @@
@Test
public void testShouldDrain() {
- DrainUtils.writeDrainNotification(coordinatorStreamStore, TEST_DEPLOYMENT_ID);
+ DrainUtils.writeDrainNotification(coordinatorStreamStore, TEST_RUN_ID);
NamespaceAwareCoordinatorStreamStore drainStore =
new NamespaceAwareCoordinatorStreamStore(coordinatorStreamStore, DrainUtils.DRAIN_METADATA_STORE_NAMESPACE);
- Assert.assertTrue(DrainMonitor.shouldDrain(drainStore, TEST_DEPLOYMENT_ID));
+ Assert.assertTrue(DrainMonitor.shouldDrain(drainStore, TEST_RUN_ID));
// Cleanup old drain message
DrainUtils.cleanup(coordinatorStreamStore, CONFIG);
- final String mismatchedDeploymentId = "bar";
- DrainUtils.writeDrainNotification(coordinatorStreamStore, mismatchedDeploymentId);
- Assert.assertFalse(DrainMonitor.shouldDrain(drainStore, TEST_DEPLOYMENT_ID));
+ final String mismatchedRunId = "bar";
+ DrainUtils.writeDrainNotification(coordinatorStreamStore, mismatchedRunId);
+ Assert.assertFalse(DrainMonitor.shouldDrain(drainStore, TEST_RUN_ID));
}
}
diff --git a/samza-core/src/test/java/org/apache/samza/drain/DrainUtilsTests.java b/samza-core/src/test/java/org/apache/samza/drain/DrainUtilsTests.java
index 1726535..43b91ad 100644
--- a/samza-core/src/test/java/org/apache/samza/drain/DrainUtilsTests.java
+++ b/samza-core/src/test/java/org/apache/samza/drain/DrainUtilsTests.java
@@ -45,11 +45,11 @@
* Tests for {@link DrainUtils}
* */
public class DrainUtilsTests {
- private static final String TEST_DEPLOYMENT_ID = "foo";
+ private static final String TEST_RUN_ID = "foo";
private static final Config CONFIG = new MapConfig(ImmutableMap.of(
"job.name", "test-job",
"job.coordinator.system", "test-kafka",
- ApplicationConfig.APP_RUN_ID, TEST_DEPLOYMENT_ID));
+ ApplicationConfig.APP_RUN_ID, TEST_RUN_ID));
private CoordinatorStreamStore coordinatorStreamStore;
@@ -68,17 +68,17 @@
@Test
public void testWrites() {
- String deploymentId1 = "foo1";
- String deploymentId2 = "foo2";
- String deploymentId3 = "foo3";
+ String runId1 = "foo1";
+ String runId2 = "foo2";
+ String runId3 = "foo3";
- UUID uuid1 = DrainUtils.writeDrainNotification(coordinatorStreamStore, deploymentId1);
- UUID uuid2 = DrainUtils.writeDrainNotification(coordinatorStreamStore, deploymentId2);
- UUID uuid3 = DrainUtils.writeDrainNotification(coordinatorStreamStore, deploymentId3);
+ UUID uuid1 = DrainUtils.writeDrainNotification(coordinatorStreamStore, runId1);
+ UUID uuid2 = DrainUtils.writeDrainNotification(coordinatorStreamStore, runId2);
+ UUID uuid3 = DrainUtils.writeDrainNotification(coordinatorStreamStore, runId3);
- DrainNotification expectedDrainNotification1 = new DrainNotification(uuid1, deploymentId1);
- DrainNotification expectedDrainNotification2 = new DrainNotification(uuid2, deploymentId2);
- DrainNotification expectedDrainNotification3 = new DrainNotification(uuid3, deploymentId3);
+ DrainNotification expectedDrainNotification1 = new DrainNotification(uuid1, runId1);
+ DrainNotification expectedDrainNotification2 = new DrainNotification(uuid2, runId2);
+ DrainNotification expectedDrainNotification3 = new DrainNotification(uuid3, runId3);
Set<DrainNotification> expectedDrainNotifications = new HashSet<>(Arrays.asList(expectedDrainNotification1,
expectedDrainNotification2, expectedDrainNotification3));
@@ -90,23 +90,23 @@
@Test
public void testCleanup() {
- DrainUtils.writeDrainNotification(coordinatorStreamStore, TEST_DEPLOYMENT_ID);
+ DrainUtils.writeDrainNotification(coordinatorStreamStore, TEST_RUN_ID);
DrainUtils.cleanup(coordinatorStreamStore, CONFIG);
final Optional<List<DrainNotification>> drainNotifications1 = readDrainNotificationMessages(coordinatorStreamStore);
Assert.assertFalse(drainNotifications1.isPresent());
- final String deploymentId = "bar";
- DrainUtils.writeDrainNotification(coordinatorStreamStore, deploymentId);
+ final String runId = "bar";
+ DrainUtils.writeDrainNotification(coordinatorStreamStore, runId);
DrainUtils.cleanup(coordinatorStreamStore, CONFIG);
final Optional<List<DrainNotification>> drainNotifications2 = readDrainNotificationMessages(coordinatorStreamStore);
Assert.assertTrue(drainNotifications2.isPresent());
- Assert.assertEquals(deploymentId, drainNotifications2.get().get(0).getDeploymentId());
+ Assert.assertEquals(runId, drainNotifications2.get().get(0).getRunId());
}
@Test
public void testCleanupAll() {
- DrainUtils.writeDrainNotification(coordinatorStreamStore, TEST_DEPLOYMENT_ID);
- DrainUtils.writeDrainNotification(coordinatorStreamStore, TEST_DEPLOYMENT_ID);
+ DrainUtils.writeDrainNotification(coordinatorStreamStore, TEST_RUN_ID);
+ DrainUtils.writeDrainNotification(coordinatorStreamStore, TEST_RUN_ID);
DrainUtils.writeDrainNotification(coordinatorStreamStore, "bar");
DrainUtils.cleanupAll(coordinatorStreamStore);
final Optional<List<DrainNotification>> drainNotifications = readDrainNotificationMessages(coordinatorStreamStore);
diff --git a/samza-core/src/test/scala/org/apache/samza/container/TestSamzaContainer.scala b/samza-core/src/test/scala/org/apache/samza/container/TestSamzaContainer.scala
index d9bb916..6b2154e 100644
--- a/samza-core/src/test/scala/org/apache/samza/container/TestSamzaContainer.scala
+++ b/samza-core/src/test/scala/org/apache/samza/container/TestSamzaContainer.scala
@@ -57,7 +57,7 @@
@Mock
private var taskInstance: TaskInstance = null
@Mock
- private var runLoop: Runnable = null
+ private var runLoop: RunLoop = null
@Mock
private var systemAdmins: SystemAdmins = null
@Mock
diff --git a/samza-test/src/main/java/org/apache/samza/test/framework/TestRunner.java b/samza-test/src/main/java/org/apache/samza/test/framework/TestRunner.java
index 943a8ca..60323d6 100644
--- a/samza-test/src/main/java/org/apache/samza/test/framework/TestRunner.java
+++ b/samza-test/src/main/java/org/apache/samza/test/framework/TestRunner.java
@@ -20,6 +20,7 @@
package org.apache.samza.test.framework;
import com.google.common.base.Preconditions;
+import com.google.common.collect.Lists;
import java.io.File;
import java.time.Duration;
import java.util.HashMap;
@@ -28,6 +29,9 @@
import java.util.List;
import java.util.Map;
import java.util.Set;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;
import org.apache.commons.lang3.RandomStringUtils;
import org.apache.samza.SamzaException;
@@ -93,10 +97,12 @@
private static final Logger LOG = LoggerFactory.getLogger(TestRunner.class);
private static final String JOB_DEFAULT_SYSTEM = "default-samza-system";
private static final String APP_NAME = "samza-test";
+ private static final long DEFAULT_DELAY_BETWEEN_MESSAGES = 0L;
private Map<String, String> configs;
private SamzaApplication app;
private ExternalContext externalContext;
+ private InMemoryMetadataStoreFactory inMemoryMetadataStoreFactory;
/*
* inMemoryScope is a unique global key per TestRunner, this key when configured with {@link InMemorySystemDescriptor}
* provides an isolated state to run with in memory system
@@ -201,6 +207,13 @@
}
/**
+ * Get configs.
+ * */
+ public Config getConfig() {
+ return new MapConfig(configs);
+ }
+
+ /**
* Passes the user provided external context to {@link LocalApplicationRunner}
*
* @param externalContext external context provided by user
@@ -213,6 +226,18 @@
}
/**
+ * Set an InMemory MetadataStoreFactory to be used by {@link LocalApplicationRunner}
+ *
+ * @param inMemoryMetadataStoreFactory InMemoryMetadataStoreFactory
+ * @return this {@link TestRunner}
+ */
+ public TestRunner setInMemoryMetadataFactory(InMemoryMetadataStoreFactory inMemoryMetadataStoreFactory) {
+ Preconditions.checkNotNull(inMemoryMetadataStoreFactory);
+ this.inMemoryMetadataStoreFactory = inMemoryMetadataStoreFactory;
+ return this;
+ }
+
+ /**
* Adds the provided input stream with mock data to the test application.
*
* @param descriptor describes the stream that is supposed to be input to Samza application
@@ -224,10 +249,29 @@
*/
public <StreamMessageType> TestRunner addInputStream(InMemoryInputDescriptor descriptor,
List<StreamMessageType> messages) {
+ return addInputStream(descriptor, messages, DEFAULT_DELAY_BETWEEN_MESSAGES);
+ }
+
+ /**
+ * Adds the provided input stream with mock data to the test application.The mock messages will be added one at a time
+ * with a delay between messages instead of adding all at once.
+ * Default configs and user added configs have
+ * a higher precedence over system and stream descriptor generated configs.
+ * @param descriptor describes the stream that is supposed to be input to Samza application
+ * @param messages map whose key is partitionId and value is messages in the partition. These message should always
+ * be deserialized
+ * @param delayBetweenMessagesInMillis delay between messages
+ * @param <StreamMessageType> message with null key or a KV {@link org.apache.samza.operators.KV}.
+ * A key of which represents key of {@link org.apache.samza.system.IncomingMessageEnvelope} or
+ * {@link org.apache.samza.system.OutgoingMessageEnvelope} and value is message
+ * @return this {@link TestRunner}
+ */
+ public <StreamMessageType> TestRunner addInputStream(InMemoryInputDescriptor descriptor,
+ List<StreamMessageType> messages, long delayBetweenMessagesInMillis) {
Preconditions.checkNotNull(descriptor, messages);
- Map<Integer, Iterable<StreamMessageType>> partitionData = new HashMap<Integer, Iterable<StreamMessageType>>();
+ Map<Integer, Iterable<StreamMessageType>> partitionData = new HashMap<>();
partitionData.put(0, messages);
- initializeInMemoryInputStream(descriptor, partitionData);
+ initializeInMemoryInputStream(descriptor, partitionData, delayBetweenMessagesInMillis);
return this;
}
@@ -244,10 +288,28 @@
*/
public <StreamMessageType> TestRunner addInputStream(InMemoryInputDescriptor descriptor,
Map<Integer, ? extends Iterable<StreamMessageType>> messages) {
+ return addInputStream(descriptor, messages, DEFAULT_DELAY_BETWEEN_MESSAGES);
+ }
+
+ /**
+ * Adds the provided input stream with mock data to the test application. The mock messages will be added one at a time
+ * with a delay between messages instead of adding all at once. TestRunner will cycle through all partitions once to add
+ * a message each per partition. This will be repeated periodically till all messages are exhausted.
+ * Default configs and user added configs have a higher precedence over system and stream descriptor generated configs.
+ * @param descriptor describes the stream that is supposed to be input to Samza application
+ * @param messages map whose key is partitionId and value is messages in the partition. These message should always
+ * be deserialized
+ * @param delayBetweenMessagesInMillis delay bewtween messages.
+ * @param <StreamMessageType> message with null key or a KV {@link org.apache.samza.operators.KV}.
+ * A key of which represents key of {@link org.apache.samza.system.IncomingMessageEnvelope} or
+ * {@link org.apache.samza.system.OutgoingMessageEnvelope} and value is message
+ * @return this {@link TestRunner}
+ */
+ public <StreamMessageType> TestRunner addInputStream(InMemoryInputDescriptor descriptor,
+ Map<Integer, ? extends Iterable<StreamMessageType>> messages, long delayBetweenMessagesInMillis) {
Preconditions.checkNotNull(descriptor, messages);
- Map<Integer, Iterable<StreamMessageType>> partitionData = new HashMap<Integer, Iterable<StreamMessageType>>();
- partitionData.putAll(messages);
- initializeInMemoryInputStream(descriptor, partitionData);
+ Map<Integer, Iterable<StreamMessageType>> partitionData = new HashMap<>(messages);
+ initializeInMemoryInputStream(descriptor, partitionData, delayBetweenMessagesInMillis);
return this;
}
@@ -291,7 +353,10 @@
// Cleaning store directories to ensure current run does not pick up state from previous run
deleteStoreDirectories();
Config config = new MapConfig(JobPlanner.generateSingleJobConfig(configs));
- final LocalApplicationRunner runner = new LocalApplicationRunner(app, config, new InMemoryMetadataStoreFactory());
+ InMemoryMetadataStoreFactory metadataStoreFactory = inMemoryMetadataStoreFactory != null
+ ? inMemoryMetadataStoreFactory
+ : new InMemoryMetadataStoreFactory();
+ final LocalApplicationRunner runner = new LocalApplicationRunner(app, config, metadataStoreFactory);
runner.run(externalContext);
if (!runner.waitForFinish(timeout)) {
throw new SamzaException("Timed out waiting for application to finish");
@@ -378,7 +443,7 @@
* @param descriptor describes a stream to initialize with the in memory system
*/
private <StreamMessageType> void initializeInMemoryInputStream(InMemoryInputDescriptor<?> descriptor,
- Map<Integer, Iterable<StreamMessageType>> partitionData) {
+ Map<Integer, Iterable<StreamMessageType>> partitionData, long delayInMillis) {
String systemName = descriptor.getSystemName();
String streamName = (String) descriptor.getPhysicalName().orElse(descriptor.getStreamId());
if (this.app instanceof LegacyTaskApplication) {
@@ -402,19 +467,64 @@
factory.getAdmin(systemName, config).createStream(spec);
InMemorySystemProducer producer = (InMemorySystemProducer) factory.getProducer(systemName, config, null);
SystemStream sysStream = new SystemStream(systemName, streamName);
- partitionData.forEach((partitionId, partition) -> {
- partition.forEach(e -> {
- Object key = e instanceof KV ? ((KV) e).getKey() : null;
- Object value = e instanceof KV ? ((KV) e).getValue() : e;
- if (value instanceof IncomingMessageEnvelope) {
- producer.send((IncomingMessageEnvelope) value);
- } else {
- producer.send(systemName, new OutgoingMessageEnvelope(sysStream, Integer.valueOf(partitionId), key, value));
- }
+ if (delayInMillis > 0) {
+ delayedInitialization(partitionData, producer, sysStream, delayInMillis);
+ } else {
+ partitionData.forEach((partitionId, partition) -> {
+ partition.forEach(e -> {
+ Object key = e instanceof KV ? ((KV) e).getKey() : null;
+ Object value = e instanceof KV ? ((KV) e).getValue() : e;
+ if (value instanceof IncomingMessageEnvelope) {
+ producer.send((IncomingMessageEnvelope) value);
+ } else {
+ producer.send(systemName, new OutgoingMessageEnvelope(sysStream, Integer.valueOf(partitionId), key, value));
+ }
+ });
+ producer.send(systemName, new OutgoingMessageEnvelope(sysStream, Integer.valueOf(partitionId), null,
+ new EndOfStreamMessage(null)));
});
- producer.send(systemName, new OutgoingMessageEnvelope(sysStream, Integer.valueOf(partitionId), null,
- new EndOfStreamMessage(null)));
- });
+ }
+ }
+
+ private <StreamMessageType> void delayedInitialization(Map<Integer, Iterable<StreamMessageType>> partitionData,
+ InMemorySystemProducer producer, SystemStream systemStream, long delayInMillis) {
+ final Set<Integer> endOfStreamPartitions = new HashSet<>();
+ final int numPartitions = partitionData.size();
+ Map<Integer, LinkedList<StreamMessageType>> messageQueuesByPartition = partitionData.entrySet()
+ .stream()
+ .collect(Collectors.toMap(
+ Map.Entry::getKey,
+ entry -> Lists.newLinkedList(entry.getValue())));
+
+ final ScheduledExecutorService executor = Executors.newSingleThreadScheduledExecutor();
+ executor.scheduleAtFixedRate(new Runnable() {
+ @Override
+ public void run() {
+ if (endOfStreamPartitions.size() == numPartitions) {
+ // shutdown scheduled executor as EOS is reached for all partitions
+ executor.shutdownNow();
+ } else {
+ messageQueuesByPartition.forEach((partitionId, messages) -> {
+ if (messages.isEmpty() && !endOfStreamPartitions.contains(partitionId)) {
+ // if end of input is reached, send EOS for the partition and add the partition to the EOS partition set
+ // to ensure we don't send EOS again for the partition
+ producer.send(systemStream.getSystem(),
+ new OutgoingMessageEnvelope(systemStream, Integer.valueOf(partitionId), null, new EndOfStreamMessage(null)));
+ endOfStreamPartitions.add(partitionId);
+ } else if (!messages.isEmpty()) {
+ final StreamMessageType e = messageQueuesByPartition.get(partitionId).removeFirst();
+ final Object key = e instanceof KV ? ((KV) e).getKey() : null;
+ final Object value = e instanceof KV ? ((KV) e).getValue() : e;
+ if (value instanceof IncomingMessageEnvelope) {
+ producer.send((IncomingMessageEnvelope) value);
+ } else {
+ producer.send(systemStream.getSystem(), new OutgoingMessageEnvelope(systemStream, Integer.valueOf(partitionId), key, value));
+ }
+ }
+ });
+ }
+ }
+ }, 0L, delayInMillis, TimeUnit.MILLISECONDS);
}
private void deleteStoreDirectories() {
diff --git a/samza-test/src/test/java/org/apache/samza/test/controlmessages/TestData.java b/samza-test/src/test/java/org/apache/samza/test/TestData.java
similarity index 97%
rename from samza-test/src/test/java/org/apache/samza/test/controlmessages/TestData.java
rename to samza-test/src/test/java/org/apache/samza/test/TestData.java
index 52d8380..5b2f137 100644
--- a/samza-test/src/test/java/org/apache/samza/test/controlmessages/TestData.java
+++ b/samza-test/src/test/java/org/apache/samza/test/TestData.java
@@ -17,7 +17,7 @@
* under the License.
*/
-package org.apache.samza.test.controlmessages;
+package org.apache.samza.test;
import java.io.Serializable;
import org.apache.samza.SamzaException;
diff --git a/samza-test/src/test/java/org/apache/samza/test/controlmessages/WatermarkIntegrationTest.java b/samza-test/src/test/java/org/apache/samza/test/controlmessages/WatermarkIntegrationTest.java
index 849f64c..bbe03ae 100644
--- a/samza-test/src/test/java/org/apache/samza/test/controlmessages/WatermarkIntegrationTest.java
+++ b/samza-test/src/test/java/org/apache/samza/test/controlmessages/WatermarkIntegrationTest.java
@@ -67,8 +67,8 @@
import org.apache.samza.system.SystemStreamPartition;
import org.apache.samza.task.StreamOperatorTask;
import org.apache.samza.task.TestStreamOperatorTask;
-import org.apache.samza.test.controlmessages.TestData.PageView;
-import org.apache.samza.test.controlmessages.TestData.PageViewJsonSerdeFactory;
+import org.apache.samza.test.TestData.PageView;
+import org.apache.samza.test.TestData.PageViewJsonSerdeFactory;
import org.apache.samza.test.harness.IntegrationTestHarness;
import org.apache.samza.test.util.SimpleSystemAdmin;
import org.apache.samza.test.util.TestStreamConsumer;
diff --git a/samza-test/src/test/java/org/apache/samza/test/drain/DrainHighLevelApiIntegrationTest.java b/samza-test/src/test/java/org/apache/samza/test/drain/DrainHighLevelApiIntegrationTest.java
new file mode 100644
index 0000000..c1a6f91
--- /dev/null
+++ b/samza-test/src/test/java/org/apache/samza/test/drain/DrainHighLevelApiIntegrationTest.java
@@ -0,0 +1,127 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.samza.test.drain;
+
+import com.google.common.collect.ImmutableMap;
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.UUID;
+import java.util.concurrent.Callable;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import org.apache.samza.application.StreamApplication;
+import org.apache.samza.application.descriptors.StreamApplicationDescriptor;
+import org.apache.samza.config.ApplicationConfig;
+import org.apache.samza.config.Config;
+import org.apache.samza.config.JobConfig;
+import org.apache.samza.drain.DrainUtils;
+import org.apache.samza.metadatastore.InMemoryMetadataStoreFactory;
+import org.apache.samza.metadatastore.MetadataStore;
+import org.apache.samza.metrics.MetricsRegistryMap;
+import org.apache.samza.operators.KV;
+import org.apache.samza.serializers.IntegerSerde;
+import org.apache.samza.serializers.KVSerde;
+import org.apache.samza.serializers.NoOpSerde;
+import org.apache.samza.system.descriptors.DelegatingSystemDescriptor;
+import org.apache.samza.system.descriptors.GenericInputDescriptor;
+import org.apache.samza.test.framework.TestRunner;
+import org.apache.samza.test.framework.system.descriptors.InMemoryInputDescriptor;
+import org.apache.samza.test.framework.system.descriptors.InMemorySystemDescriptor;
+import org.apache.samza.test.table.TestTableData;
+import org.apache.samza.test.table.TestTableData.PageView;
+import org.junit.Ignore;
+import org.junit.Test;
+
+import static org.junit.Assert.*;
+
+
+/**
+ * End to end integration test to check drain functionality with samza high-level API.
+ */
+public class DrainHighLevelApiIntegrationTest {
+ private static final List<PageView> RECEIVED = new ArrayList<>();
+ private static final String SYSTEM_NAME = "test";
+ private static final String STREAM_ID = "PageView";
+
+ private static class PageViewEventCountHighLevelApplication implements StreamApplication {
+ @Override
+ public void describe(StreamApplicationDescriptor appDescriptor) {
+ DelegatingSystemDescriptor sd = new DelegatingSystemDescriptor(SYSTEM_NAME);
+ GenericInputDescriptor<KV<String, PageView>> isd =
+ sd.getInputDescriptor(STREAM_ID, KVSerde.of(new NoOpSerde<>(), new NoOpSerde<>()));
+ appDescriptor.getInputStream(isd)
+ .map(KV::getValue)
+ .partitionBy(PageView::getMemberId, pv -> pv,
+ KVSerde.of(new IntegerSerde(), new TestTableData.PageViewJsonSerde()), "p1")
+ .sink((m, collector, coordinator) -> {
+ RECEIVED.add(m.getValue());
+ });
+ }
+ }
+
+ // The test can be occasionally flaky, so we set Ignore annotation
+ // Remove ignore annotation and run the test as follows:
+ // ./gradlew :samza-test:test --tests org.apache.samza.test.drain.DrainHighLevelApiIntegrationTest -PscalaSuffix=2.12
+ @Ignore
+ @Test
+ public void testPipeline() {
+ String runId = "DrainTestId";
+ int numPageViews = 40;
+
+ InMemorySystemDescriptor isd = new InMemorySystemDescriptor(SYSTEM_NAME);
+ InMemoryInputDescriptor<TestTableData.PageView> inputDescriptor = isd.getInputDescriptor(STREAM_ID, new NoOpSerde<>());
+ InMemoryMetadataStoreFactory metadataStoreFactory = new InMemoryMetadataStoreFactory();
+
+ Map<String, String> customConfig = ImmutableMap.of(
+ ApplicationConfig.APP_RUN_ID, runId,
+ JobConfig.DRAIN_MONITOR_POLL_INTERVAL_MILLIS, "100",
+ JobConfig.DRAIN_MONITOR_ENABLED, "true");
+
+ // Create a TestRunner
+ // Set a InMemoryMetadataFactory. We will use this factory in the test to create a metadata store and
+ // write drain message to it
+ // Mock data comprises of 40 messages across 4 partitions. TestRunner adds a 1 second delay between messages
+ // per partition when writing messages to the InputStream
+ TestRunner testRunner = TestRunner.of(new PageViewEventCountHighLevelApplication())
+ .setInMemoryMetadataFactory(metadataStoreFactory)
+ .addConfig(customConfig)
+ .addInputStream(inputDescriptor, TestTableData.generatePartitionedPageViews(numPageViews, 4), 1000L);
+
+ Config configFromRunner = testRunner.getConfig();
+ MetadataStore metadataStore = metadataStoreFactory.getMetadataStore("NoOp", configFromRunner, new MetricsRegistryMap());
+
+ // write drain message after a delay
+ ScheduledExecutorService executorService = Executors.newSingleThreadScheduledExecutor();
+ executorService.schedule(new Callable<String>() {
+ @Override
+ public String call() throws Exception {
+ UUID uuid = DrainUtils.writeDrainNotification(metadataStore, runId);
+ return uuid.toString();
+ }
+ }, 2000L, TimeUnit.MILLISECONDS);
+
+ testRunner.run(Duration.ofSeconds(20));
+
+ assertTrue(RECEIVED.size() < numPageViews && RECEIVED.size() > 0);
+ }
+}
diff --git a/samza-test/src/test/java/org/apache/samza/test/drain/DrainLowLevelApiIntegrationTest.java b/samza-test/src/test/java/org/apache/samza/test/drain/DrainLowLevelApiIntegrationTest.java
new file mode 100644
index 0000000..7579c6e
--- /dev/null
+++ b/samza-test/src/test/java/org/apache/samza/test/drain/DrainLowLevelApiIntegrationTest.java
@@ -0,0 +1,156 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.samza.test.drain;
+
+import com.google.common.collect.ImmutableMap;
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.UUID;
+import java.util.concurrent.Callable;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import org.apache.samza.application.TaskApplication;
+import org.apache.samza.application.descriptors.TaskApplicationDescriptor;
+import org.apache.samza.config.ApplicationConfig;
+import org.apache.samza.config.Config;
+import org.apache.samza.config.JobConfig;
+import org.apache.samza.context.Context;
+import org.apache.samza.drain.DrainUtils;
+import org.apache.samza.metadatastore.InMemoryMetadataStoreFactory;
+import org.apache.samza.metadatastore.MetadataStore;
+import org.apache.samza.metrics.MetricsRegistryMap;
+import org.apache.samza.serializers.NoOpSerde;
+import org.apache.samza.system.IncomingMessageEnvelope;
+import org.apache.samza.system.descriptors.DelegatingSystemDescriptor;
+import org.apache.samza.system.descriptors.GenericInputDescriptor;
+import org.apache.samza.task.DrainListenerTask;
+import org.apache.samza.task.EndOfStreamListenerTask;
+import org.apache.samza.task.InitableTask;
+import org.apache.samza.task.MessageCollector;
+import org.apache.samza.task.StreamTask;
+import org.apache.samza.task.StreamTaskFactory;
+import org.apache.samza.task.TaskCoordinator;
+import org.apache.samza.test.framework.TestRunner;
+import org.apache.samza.test.framework.system.descriptors.InMemoryInputDescriptor;
+import org.apache.samza.test.framework.system.descriptors.InMemorySystemDescriptor;
+import org.apache.samza.test.table.TestTableData;
+import org.junit.Ignore;
+import org.junit.Test;
+
+import static org.junit.Assert.*;
+
+
+/**
+ * End to end integration test to check drain functionality with samza low-level API.
+ * */
+public class DrainLowLevelApiIntegrationTest {
+ private static final List<TestTableData.PageView> RECEIVED = new ArrayList<>();
+
+ private static Integer drainCounter = 0;
+ private static Integer eosCounter = 0;
+
+ private static final String SYSTEM_NAME = "test";
+ private static final String STREAM_ID = "PageView";
+
+ private static class PageViewEventCountLowLevelApplication implements TaskApplication {
+ @Override
+ public void describe(TaskApplicationDescriptor appDescriptor) {
+ DelegatingSystemDescriptor ksd = new DelegatingSystemDescriptor(SYSTEM_NAME);
+ GenericInputDescriptor<TestTableData.PageView> pageViewIsd = ksd.getInputDescriptor(STREAM_ID, new NoOpSerde<>());
+ appDescriptor
+ .withInputStream(pageViewIsd)
+ .withTaskFactory((StreamTaskFactory) PageViewEventCountStreamTask::new);
+ }
+ }
+
+ private static class PageViewEventCountStreamTask implements StreamTask, InitableTask, DrainListenerTask, EndOfStreamListenerTask {
+ public PageViewEventCountStreamTask() {
+ }
+
+ @Override
+ public void init(Context context) {
+ }
+
+ @Override
+ public void process(IncomingMessageEnvelope message, MessageCollector collector, TaskCoordinator coordinator) {
+ TestTableData.PageView pv = (TestTableData.PageView) message.getMessage();
+ RECEIVED.add(pv);
+ }
+
+ @Override
+ public void onDrain(MessageCollector collector, TaskCoordinator coordinator) throws Exception {
+ drainCounter++;
+ }
+
+ @Override
+ public void onEndOfStream(MessageCollector collector, TaskCoordinator coordinator) throws Exception {
+ eosCounter++;
+ }
+ }
+
+ // The test can be occasionally flaky, so we set Ignore annotation
+ // Remove ignore annotation and run the test as follows:
+ // ./gradlew :samza-test:test --tests org.apache.samza.test.drain.DrainHighLevelApiIntegrationTest -PscalaSuffix=2.12
+ @Ignore
+ @Test
+ public void testPipeline() {
+ int numPageViews = 40;
+ InMemorySystemDescriptor isd = new InMemorySystemDescriptor("test");
+ InMemoryInputDescriptor<TestTableData.PageView> inputDescriptor = isd.getInputDescriptor("PageView", new NoOpSerde<>());
+ InMemoryMetadataStoreFactory metadataStoreFactory = new InMemoryMetadataStoreFactory();
+
+ String runId = "DrainTestId";
+ Map<String, String> customConfig = ImmutableMap.of(
+ ApplicationConfig.APP_RUN_ID, runId,
+ JobConfig.DRAIN_MONITOR_POLL_INTERVAL_MILLIS, "100",
+ JobConfig.DRAIN_MONITOR_ENABLED, "true");
+
+ // Create a TestRunner
+ // Set a InMemoryMetadataFactory. We will use this factory in the test to create a metadata store and
+ // write drain message to it
+ // Mock data comprises of 40 messages across 4 partitions. TestRunner adds a 1 second delay between messages
+ // per partition when writing messages to the InputStream
+ TestRunner testRunner = TestRunner.of(new PageViewEventCountLowLevelApplication())
+ .setInMemoryMetadataFactory(metadataStoreFactory)
+ .addConfig(customConfig)
+ .addInputStream(inputDescriptor, TestTableData.generatePartitionedPageViews(numPageViews, 4), 1000L);
+
+ Config configFromRunner = testRunner.getConfig();
+ MetadataStore
+ metadataStore = metadataStoreFactory.getMetadataStore("NoOp", configFromRunner, new MetricsRegistryMap());
+
+ // write drain message after a delay
+ ScheduledExecutorService executorService = Executors.newSingleThreadScheduledExecutor();
+ executorService.schedule(new Callable<String>() {
+ @Override
+ public String call() throws Exception {
+ UUID uuid = DrainUtils.writeDrainNotification(metadataStore, runId);
+ return uuid.toString();
+ }
+ }, 2000L, TimeUnit.MILLISECONDS);
+
+ testRunner.run(Duration.ofSeconds(25));
+
+ assertTrue(RECEIVED.size() < numPageViews && RECEIVED.size() > 0);
+ }
+}
diff --git a/samza-test/src/test/java/org/apache/samza/test/framework/StreamApplicationIntegrationTest.java b/samza-test/src/test/java/org/apache/samza/test/framework/StreamApplicationIntegrationTest.java
index 6afc77c..2a6e70c 100644
--- a/samza-test/src/test/java/org/apache/samza/test/framework/StreamApplicationIntegrationTest.java
+++ b/samza-test/src/test/java/org/apache/samza/test/framework/StreamApplicationIntegrationTest.java
@@ -42,7 +42,7 @@
import org.apache.samza.system.kafka.descriptors.KafkaOutputDescriptor;
import org.apache.samza.system.kafka.descriptors.KafkaSystemDescriptor;
import org.apache.samza.table.Table;
-import org.apache.samza.test.controlmessages.TestData;
+import org.apache.samza.test.TestData;
import org.apache.samza.test.framework.system.descriptors.InMemoryInputDescriptor;
import org.apache.samza.test.framework.system.descriptors.InMemoryOutputDescriptor;
import org.apache.samza.test.framework.system.descriptors.InMemorySystemDescriptor;
@@ -53,7 +53,7 @@
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
-import static org.apache.samza.test.controlmessages.TestData.PageView;
+import static org.apache.samza.test.TestData.PageView;
public class StreamApplicationIntegrationTest {
private static final Logger LOG = LoggerFactory.getLogger(StreamApplicationIntegrationTest.class);
