stringBufferMapEntry : stringBufferMap.entrySet()) {
- AppendKey key = stringBufferMapEntry.getKey();
- StringBuilder buffer = stringBufferMapEntry.getValue();
- sizeProcessed += buffer.length();
- try {
- listener.onParagraphOutputAppend(key.noteId, key.paragraphId, key.index,
- key.executionOwner, buffer.toString());
- } catch (RuntimeException e) {
- // One stale append must not abort another paragraph's synchronous drain.
- LOGGER.warn("Failed to append output for {}", key, e);
- }
+ private long flush(ParagraphOutputKey key, StringBuilder data) {
+ long size = data.length();
+ try {
+ listener.onParagraphOutputAppend(
+ key.noteId, key.paragraphId, key.index, key.executionOwner, data.toString());
+ } catch (RuntimeException e) {
+ // A stale paragraph must not abort delivery of later output in this drain.
+ LOGGER.warn("Failed to append output for note {} paragraph {}",
+ key.noteId, key.paragraphId, e);
}
- stringBufferMap.clear();
- return sizeProcessed;
+ data.setLength(0);
+ return size;
}
- /**
- * Identifies one stream of appended output. An owner name can contain any character, so the
- * parts are kept separate instead of being joined into a delimited string.
- */
- private static final class AppendKey {
+ private static final class ParagraphOutputKey {
private final String noteId;
private final String paragraphId;
private final int index;
private final String executionOwner;
- private AppendKey(String noteId, String paragraphId, int index, String executionOwner) {
- this.noteId = noteId;
- this.paragraphId = paragraphId;
- this.index = index;
- this.executionOwner = executionOwner;
+ private ParagraphOutputKey(AppendOutputBuffer append) {
+ noteId = append.getNoteId();
+ paragraphId = append.getParagraphId();
+ index = append.getIndex();
+ executionOwner = append.getExecutionOwner();
+ }
+
+ private boolean matches(AppendOutputBuffer append) {
+ return index == append.getIndex()
+ && Objects.equals(noteId, append.getNoteId())
+ && Objects.equals(paragraphId, append.getParagraphId())
+ && Objects.equals(executionOwner, append.getExecutionOwner());
}
@Override
- public boolean equals(Object o) {
- if (this == o) {
+ public boolean equals(Object other) {
+ if (this == other) {
return true;
}
- if (!(o instanceof AppendKey)) {
+ if (!(other instanceof ParagraphOutputKey)) {
return false;
}
- AppendKey other = (AppendKey) o;
- return index == other.index
- && Objects.equals(noteId, other.noteId)
- && Objects.equals(paragraphId, other.paragraphId)
- && Objects.equals(executionOwner, other.executionOwner);
+ ParagraphOutputKey key = (ParagraphOutputKey) other;
+ return index == key.index && Objects.equals(noteId, key.noteId)
+ && Objects.equals(paragraphId, key.paragraphId)
+ && Objects.equals(executionOwner, key.executionOwner);
}
@Override
public int hashCode() {
- return Objects.hash(noteId, paragraphId, index, executionOwner);
- }
-
- @Override
- public String toString() {
- return "note " + noteId + " paragraph " + paragraphId + " index " + index
- + " executionOwner " + executionOwner;
+ int hash = Objects.hashCode(noteId);
+ hash = 31 * hash + Objects.hashCode(paragraphId);
+ hash = 31 * hash + index;
+ return 31 * hash + Objects.hashCode(executionOwner);
}
}
-
- public void appendBuffer(String noteId, String paragraphId, int index, String executionOwner,
- String outputToAppend) {
- queue.offer(
- new AppendOutputBuffer(noteId, paragraphId, index, executionOwner, outputToAppend));
- }
-
- /** Enqueues a replacement; callers needing completion must also invoke run(). */
- public void updateBuffer(String noteId, String paragraphId, int index, String executionOwner,
- InterpreterResult.Type type, String output) {
- queue.offer(
- new UpdateOutputBuffer(noteId, paragraphId, index, executionOwner, type, output));
- }
}
diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/remote/ParagraphOutputDispatcher.java b/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/remote/ParagraphOutputDispatcher.java
new file mode 100644
index 00000000000..0d2253a1ca8
--- /dev/null
+++ b/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/remote/ParagraphOutputDispatcher.java
@@ -0,0 +1,453 @@
+/*
+ * 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.zeppelin.interpreter.remote;
+
+import static org.apache.zeppelin.conf.ZeppelinConfiguration.ConfVars
+ .ZEPPELIN_INTERPRETER_OUTPUT_EVENTS_PER_BATCH;
+import static org.apache.zeppelin.conf.ZeppelinConfiguration.ConfVars
+ .ZEPPELIN_INTERPRETER_OUTPUT_WORKER_COUNT;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import java.util.ArrayDeque;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.FutureTask;
+import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.RejectedExecutionException;
+import java.util.concurrent.TimeUnit;
+import org.apache.zeppelin.interpreter.InterpreterResult;
+import org.apache.zeppelin.interpreter.InterpreterResultMessage;
+
+/**
+ * Buffers output in per-note FIFO queues, consumed by a fixed pool of workers. Only one operation
+ * owns a note at a time. Callers periodically invoke {@link #flush()} to deliver buffered appends;
+ * output boundaries request immediate delivery and expose callback completion to the caller.
+ * Checkpoints pause their note until a separate worker pool of the same configured size finishes
+ * the save callback.
+ * Queues are retired as soon as their pending and in-flight output has been processed.
+ *
+ * Accepted boundaries cannot be cancelled. Their completion acknowledges callback delivery,
+ * not note idleness. RPC callers wait outside the queue monitor.
+ */
+public class ParagraphOutputDispatcher implements AutoCloseable {
+
+ private static final Logger LOGGER = LoggerFactory.getLogger(ParagraphOutputDispatcher.class);
+
+ // Guarded by this: notes, mutable NoteQueue state, and worker startup.
+ private final Map notes = new HashMap<>();
+ private final BlockingQueue readyNotes = new LinkedBlockingQueue<>();
+ private final int eventsPerBatch;
+ private final ExecutorService outputExecutor;
+ private final ExecutorService checkpointExecutor;
+ private final int outputWorkerCount;
+ private final RemoteInterpreterProcessListener listener;
+ private final AppendOutputRunner appendRunner;
+ private volatile boolean closed;
+ private boolean outputWorkersStarted;
+
+ public ParagraphOutputDispatcher(RemoteInterpreterProcessListener listener) {
+ this(listener, ZEPPELIN_INTERPRETER_OUTPUT_WORKER_COUNT.getIntValue(),
+ ZEPPELIN_INTERPRETER_OUTPUT_EVENTS_PER_BATCH.getIntValue());
+ }
+
+ ParagraphOutputDispatcher(RemoteInterpreterProcessListener listener, int workerCount) {
+ this(listener, workerCount, ZEPPELIN_INTERPRETER_OUTPUT_EVENTS_PER_BATCH.getIntValue());
+ }
+
+ public ParagraphOutputDispatcher(RemoteInterpreterProcessListener listener, int workerCount,
+ int eventsPerBatch) {
+ if (workerCount < 1) {
+ throw new IllegalArgumentException("Output worker count must be positive");
+ }
+ if (eventsPerBatch < 1) {
+ throw new IllegalArgumentException("Output events per batch must be positive");
+ }
+
+ this.eventsPerBatch = eventsPerBatch;
+ this.listener = listener;
+ appendRunner = new AppendOutputRunner(listener);
+
+ outputExecutor = Executors.newFixedThreadPool(workerCount, runnable -> {
+ Thread thread = new Thread(runnable, "zeppelin-output-worker");
+ thread.setDaemon(true);
+ return thread;
+ });
+
+ // Keep save capacity bounded independently from output delivery.
+ checkpointExecutor = Executors.newFixedThreadPool(workerCount, runnable -> {
+ Thread thread = new Thread(runnable, "zeppelin-output-checkpoint");
+ thread.setDaemon(true);
+ return thread;
+ });
+ this.outputWorkerCount = workerCount;
+ }
+
+ /** Makes pending notes ready without waiting for any listener callback. */
+ public synchronized void flush() {
+ if (!closed) {
+ for (NoteQueue note : notes.values()) {
+ makeReady(note);
+ }
+ }
+ }
+
+ public void appendOutput(String noteId, String paragraphId, int index, String executionOwner,
+ String output) {
+ enqueue(noteId, OutputEvent.append(
+ new AppendOutputBuffer(noteId, paragraphId, index, executionOwner, output)), false);
+ }
+
+ public Future updateOutput(String noteId, String paragraphId, int index,
+ String executionOwner, InterpreterResult.Type type,
+ String output) {
+ return enqueueBoundary(noteId, () -> listener.onParagraphOutputUpdated(
+ noteId, paragraphId, index, executionOwner, type, output));
+ }
+
+ public Future updateAllOutput(String noteId, String paragraphId, String executionOwner,
+ List messages) {
+ // The caller may change its list before this queued operation gets a worker.
+ List replacements = new ArrayList<>(messages);
+ return enqueueBoundary(noteId, () -> {
+ // Clear and replacements must stay together; later appends belong to the replaced output.
+ listener.onParagraphOutputClear(noteId, paragraphId, executionOwner);
+ for (int i = 0; i < replacements.size(); i++) {
+ InterpreterResultMessage message = replacements.get(i);
+ listener.onParagraphOutputUpdated(noteId, paragraphId, i, executionOwner,
+ message.getType(), message.getData());
+ }
+ });
+ }
+
+ public Future checkpointOutput(String noteId, String paragraphId) {
+ return enqueueCheckpoint(noteId, () -> listener.checkpointOutput(noteId, paragraphId));
+ }
+
+ private Future enqueueBoundary(String noteId, Runnable callback) {
+ Boundary boundary = createBoundary(noteId, callback);
+ enqueue(noteId, OutputEvent.boundary(boundary), true);
+ return boundary;
+ }
+
+ private Future enqueueCheckpoint(String noteId, Runnable callback) {
+ Boundary boundary = createBoundary(noteId, callback);
+ enqueue(noteId, OutputEvent.checkpoint(boundary), true);
+ return boundary;
+ }
+
+ private Boundary createBoundary(String noteId, Runnable callback) {
+ return new Boundary(() -> {
+ long start = System.nanoTime();
+ try {
+ callback.run();
+ } catch (RuntimeException e) {
+ LOGGER.warn("Failed to process output boundary for note {}", noteId, e);
+ throw e;
+ } finally {
+ long time = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - start);
+ LOGGER.debug("Processing output boundary for note {} took {} milliseconds", noteId, time);
+ }
+ });
+ }
+
+ private synchronized void enqueue(String noteId, OutputEvent event, boolean immediate) {
+ if (closed) {
+ throw new IllegalStateException("Output dispatcher is stopped");
+ }
+
+ if (!outputWorkersStarted) {
+ outputWorkersStarted = true;
+ for (int i = 0; i < outputWorkerCount; i++) {
+ outputExecutor.execute(this::consume);
+ }
+ }
+
+ NoteQueue note = notes.computeIfAbsent(noteId, NoteQueue::new);
+ note.events.addLast(event);
+ if (immediate) {
+ makeReady(note);
+ }
+ }
+
+ // Caller must hold this monitor: the state transition and ready insertion must be atomic.
+ private void makeReady(NoteQueue note) {
+ note.flushRequested = true;
+ if (note.state == NoteState.IDLE) {
+ note.state = NoteState.READY;
+ readyNotes.offer(note);
+ }
+ }
+
+ private void consume() {
+ while (!closed && !Thread.currentThread().isInterrupted()) {
+ NoteQueue note;
+ try {
+ note = readyNotes.take();
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ return;
+ }
+
+ List batch = new ArrayList<>();
+ synchronized (this) {
+ if (closed) {
+ return;
+ }
+
+ note.state = NoteState.DELIVERING;
+ note.flushRequested = false;
+
+ while (!note.events.isEmpty() && batch.size() < eventsPerBatch) {
+ OutputEvent event = note.events.removeFirst();
+ batch.add(event);
+ if (event.kind == OutputEvent.Kind.CHECKPOINT) {
+ break;
+ }
+ }
+
+ note.inFlightBoundaries.clear();
+ for (OutputEvent event : batch) {
+ if (event.boundary != null) {
+ note.inFlightBoundaries.add(event.boundary);
+ }
+ }
+ }
+
+ boolean checkpointHandled = false;
+ boolean failed = false;
+ try {
+ checkpointHandled = deliver(note, batch);
+ } catch (Throwable t) {
+ // Fail the remaining boundaries before logging so callers cannot remain blocked.
+ // Append callbacks not yet delivered from this batch are discarded.
+ failed = true;
+ failInFlightBoundaries(note, t);
+ LOGGER.error("Failed to deliver output for note {}", note.noteId, t);
+ } finally {
+ if (!checkpointHandled) {
+ releaseNote(note, failed || batch.size() == eventsPerBatch);
+ }
+ }
+ }
+ }
+
+ /** Returns whether checkpoint processing has assumed or released ownership of the note. */
+ private boolean deliver(NoteQueue note, List batch) {
+ List appends = new ArrayList<>();
+ for (OutputEvent event : batch) {
+ if (closed) {
+ return false;
+ }
+
+ if (event.kind == OutputEvent.Kind.APPEND) {
+ appends.add(event.append);
+ } else {
+ appendRunner.run(appends, () -> !closed);
+ appends.clear();
+
+ if (closed) {
+ return false;
+ }
+
+ if (event.kind == OutputEvent.Kind.CHECKPOINT) {
+ submitCheckpoint(note, event.boundary);
+ return true;
+ }
+
+ event.boundary.run();
+ removeInFlightBoundary(note, event.boundary);
+ }
+ }
+
+ if (!closed) {
+ appendRunner.run(appends, () -> !closed);
+ }
+ return false;
+ }
+
+ private void submitCheckpoint(NoteQueue note, Boundary boundary) {
+ synchronized (this) {
+ if (closed) {
+ boundary.fail(new IllegalStateException("Output dispatcher is stopped"));
+ return;
+ }
+
+ note.inFlightBoundaries.clear();
+ note.inFlightBoundaries.add(boundary);
+ note.state = NoteState.CHECKPOINTING;
+ }
+
+ try {
+ checkpointExecutor.execute(() -> {
+ try {
+ if (!closed) {
+ boundary.run();
+ }
+ } finally {
+ releaseNote(note, true);
+ }
+ });
+ } catch (RejectedExecutionException e) {
+ boundary.fail(e);
+ releaseNote(note, true);
+ }
+ }
+
+ private synchronized void removeInFlightBoundary(NoteQueue note, Boundary boundary) {
+ note.inFlightBoundaries.remove(boundary);
+ }
+
+ private synchronized void failInFlightBoundaries(NoteQueue note, Throwable failure) {
+ for (Boundary boundary : note.inFlightBoundaries) {
+ boundary.fail(failure);
+ }
+ note.inFlightBoundaries.clear();
+ }
+
+ private synchronized void releaseNote(NoteQueue note, boolean scheduleRemainingEvents) {
+ note.inFlightBoundaries.clear();
+ note.state = NoteState.IDLE;
+
+ if (!closed) {
+ if (note.events.isEmpty()) {
+ notes.remove(note.noteId);
+ } else if (scheduleRemainingEvents || note.flushRequested) {
+ makeReady(note);
+ }
+ }
+ }
+
+ /**
+ * Stops accepting output and fails unfinished boundaries. Remaining output is discarded when
+ * workers observe shutdown; in-flight listener calls may finish. Does not wait for worker exit.
+ */
+ @Override
+ public void close() {
+ synchronized (this) {
+ if (closed) {
+ return;
+ }
+
+ closed = true;
+ IllegalStateException stopped = new IllegalStateException("Output dispatcher is stopped");
+
+ for (NoteQueue note : notes.values()) {
+ failBoundaries(note.events, stopped);
+ for (Boundary boundary : note.inFlightBoundaries) {
+ boundary.fail(stopped);
+ }
+ note.events.clear();
+ note.inFlightBoundaries.clear();
+ }
+
+ notes.clear();
+ readyNotes.clear();
+ }
+
+ outputExecutor.shutdownNow();
+ checkpointExecutor.shutdownNow();
+ }
+
+ private void failBoundaries(Iterable events, IllegalStateException stopped) {
+ for (OutputEvent event : events) {
+ if (event.boundary != null) {
+ event.boundary.fail(stopped);
+ }
+ }
+ }
+
+ // Retained note queues, not a worker termination signal.
+ synchronized int pendingNoteCount() {
+ return notes.size();
+ }
+
+ private static class NoteQueue {
+ private final String noteId;
+ private final ArrayDeque events = new ArrayDeque<>();
+ // Shutdown must release RPCs waiting on boundaries already drained from events.
+ private final List inFlightBoundaries = new ArrayList<>();
+ private NoteState state = NoteState.IDLE;
+ // A request arriving during delivery must survive until the current owner releases the note.
+ private boolean flushRequested;
+
+ private NoteQueue(String noteId) {
+ this.noteId = noteId;
+ }
+ }
+
+ private enum NoteState {
+ IDLE,
+ READY,
+ DELIVERING,
+ CHECKPOINTING
+ }
+
+ // Exactly one of append and boundary is set, according to kind.
+ private static class OutputEvent {
+ private enum Kind {
+ APPEND,
+ BOUNDARY,
+ CHECKPOINT
+ }
+
+ private final Kind kind;
+ private final AppendOutputBuffer append;
+ private final Boundary boundary;
+
+ private OutputEvent(Kind kind, AppendOutputBuffer append, Boundary boundary) {
+ this.kind = kind;
+ this.append = append;
+ this.boundary = boundary;
+ }
+
+ private static OutputEvent append(AppendOutputBuffer append) {
+ return new OutputEvent(Kind.APPEND, append, null);
+ }
+
+ private static OutputEvent boundary(Boundary boundary) {
+ return new OutputEvent(Kind.BOUNDARY, null, boundary);
+ }
+
+ private static OutputEvent checkpoint(Boundary boundary) {
+ return new OutputEvent(Kind.CHECKPOINT, null, boundary);
+ }
+ }
+
+ // FutureTask wakes get() waiters without running externally supplied completion handlers.
+ private static final class Boundary extends FutureTask {
+ private Boundary(Runnable callback) {
+ super(callback, null);
+ }
+
+ private void fail(Throwable failure) {
+ setException(failure);
+ }
+
+ @Override
+ public boolean cancel(boolean mayInterruptIfRunning) {
+ // Accepted output must stay in the FIFO even if its RPC caller stops waiting.
+ return false;
+ }
+ }
+}
diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/remote/UpdateOutputBuffer.java b/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/remote/UpdateOutputBuffer.java
deleted file mode 100644
index 1243461d177..00000000000
--- a/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/remote/UpdateOutputBuffer.java
+++ /dev/null
@@ -1,41 +0,0 @@
-/*
- * 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.zeppelin.interpreter.remote;
-
-import org.apache.zeppelin.interpreter.InterpreterResult;
-
-/**
- * This element stores the buffered update-data of paragraph's output. It shares the
- * append-data queue so that an update, which replaces a result, can never be sent
- * ahead of the appends that preceded it.
- */
-public class UpdateOutputBuffer extends AppendOutputBuffer {
-
- private final InterpreterResult.Type type;
-
- public UpdateOutputBuffer(String noteId, String paragraphId, int index, String executionOwner,
- InterpreterResult.Type type, String data) {
- super(noteId, paragraphId, index, executionOwner, data);
- this.type = type;
- }
-
- public InterpreterResult.Type getType() {
- return type;
- }
-
-}
diff --git a/zeppelin-server/src/test/java/org/apache/zeppelin/interpreter/RemoteInterpreterEventServerTest.java b/zeppelin-server/src/test/java/org/apache/zeppelin/interpreter/RemoteInterpreterEventServerTest.java
index 8fac268b9e5..df677ba9aa0 100644
--- a/zeppelin-server/src/test/java/org/apache/zeppelin/interpreter/RemoteInterpreterEventServerTest.java
+++ b/zeppelin-server/src/test/java/org/apache/zeppelin/interpreter/RemoteInterpreterEventServerTest.java
@@ -16,32 +16,42 @@
*/
package org.apache.zeppelin.interpreter;
+import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.doThrow;
import static org.mockito.Mockito.inOrder;
import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.never;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
-
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
+import org.mockito.InOrder;
import java.io.IOException;
import java.io.ObjectOutputStream;
import java.io.Serializable;
import java.lang.reflect.Field;
import java.nio.ByteBuffer;
+import java.util.Arrays;
import java.util.Collections;
+import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
-
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
import org.apache.zeppelin.conf.ZeppelinConfiguration;
-import org.apache.zeppelin.interpreter.remote.AppendOutputRunner;
+import org.apache.zeppelin.conf.ZeppelinConfiguration.ConfVars;
import org.apache.zeppelin.interpreter.remote.InvokeResourceMethodEventMessage;
+import org.apache.zeppelin.interpreter.remote.ParagraphOutputDispatcher;
import org.apache.zeppelin.interpreter.remote.RemoteInterpreterProcess;
import org.apache.zeppelin.interpreter.remote.RemoteInterpreterProcessListener;
import org.apache.zeppelin.interpreter.thrift.InterpreterRPCException;
@@ -51,8 +61,6 @@
import org.apache.zeppelin.interpreter.thrift.RemoteInterpreterResultMessage;
import org.apache.zeppelin.resource.Resource;
import org.apache.zeppelin.resource.ResourceId;
-import org.junit.jupiter.api.Test;
-import org.mockito.InOrder;
public class RemoteInterpreterEventServerTest {
@@ -60,13 +68,44 @@ public class RemoteInterpreterEventServerTest {
@Test
void updateOutputCompletesBeforeReturning() throws Exception {
RemoteInterpreterProcessListener listener = mock(RemoteInterpreterProcessListener.class);
- RemoteInterpreterEventServer server =
- serverWithRunner(listener, new AppendOutputRunner(listener));
+ RemoteInterpreterEventServer server = serverWithListener(listener);
try {
+ server.appendOutput(new OutputAppendEvent("note", "para", 0, "before", null, null));
server.updateOutput(new OutputUpdateEvent("note", "para", 0, "TEXT", "final", null, null));
- // A caller may publish terminal status as soon as the RPC returns.
- verify(listener).onParagraphOutputUpdated(
- "note", "para", 0, null, InterpreterResult.Type.TEXT, "final");
+ verify(listener).onParagraphOutputUpdated("note", "para", 0, null, InterpreterResult.Type.TEXT, "final");
+ server.appendOutput(new OutputAppendEvent("note", "para", 0, "after", null, null));
+ server.checkpointOutput("note", "para");
+ InOrder order = inOrder(listener);
+ order.verify(listener).onParagraphOutputAppend("note", "para", 0, null, "before");
+ order.verify(listener).onParagraphOutputUpdated("note", "para", 0, null, InterpreterResult.Type.TEXT, "final");
+ order.verify(listener).onParagraphOutputAppend("note", "para", 0, null, "after");
+ order.verify(listener).checkpointOutput("note", "para");
+ order.verifyNoMoreInteractions();
+ } finally {
+ server.stop();
+ }
+ }
+
+ @Test
+ void outputEventsPreserveTheirExecutionOwner() throws Exception {
+ RemoteInterpreterProcessListener listener = mock(RemoteInterpreterProcessListener.class);
+ RemoteInterpreterEventServer server = serverWithListener(listener);
+ try {
+ server.appendOutput(new OutputAppendEvent(
+ "note", "para", 0, "append", null, "alice"));
+ server.updateOutput(new OutputUpdateEvent(
+ "note", "para", 0, "TEXT", "update", null, "bob"));
+ server.updateAllOutput(new OutputUpdateAllEvent("note", "para", Collections.singletonList(
+ new RemoteInterpreterResultMessage("HTML", "replacement")), "carol"));
+
+ InOrder order = inOrder(listener);
+ order.verify(listener).onParagraphOutputAppend("note", "para", 0, "alice", "append");
+ order.verify(listener).onParagraphOutputUpdated(
+ "note", "para", 0, "bob", InterpreterResult.Type.TEXT, "update");
+ order.verify(listener).onParagraphOutputClear("note", "para", "carol");
+ order.verify(listener).onParagraphOutputUpdated(
+ "note", "para", 0, "carol", InterpreterResult.Type.HTML, "replacement");
+ order.verifyNoMoreInteractions();
} finally {
server.stop();
}
@@ -75,8 +114,7 @@ void updateOutputCompletesBeforeReturning() throws Exception {
@Test
void checkpointDrainsPendingOutputBeforeSaving() throws Exception {
RemoteInterpreterProcessListener listener = mock(RemoteInterpreterProcessListener.class);
- RemoteInterpreterEventServer server =
- serverWithRunner(listener, new AppendOutputRunner(listener));
+ RemoteInterpreterEventServer server = serverWithListener(listener);
try {
server.appendOutput(new OutputAppendEvent("note", "para", 0, "pending", null, null));
server.checkpointOutput("note", "para");
@@ -91,85 +129,361 @@ void checkpointDrainsPendingOutputBeforeSaving() throws Exception {
@Test
void updateAllIsAnOrderedClearAndReplacement() throws Exception {
RemoteInterpreterProcessListener listener = mock(RemoteInterpreterProcessListener.class);
- AppendOutputRunner runner = new AppendOutputRunner(listener);
- RemoteInterpreterEventServer server = serverWithRunner(listener, runner);
+ RemoteInterpreterEventServer server = serverWithListener(listener);
try {
- runner.appendBuffer("note", "para", 0, null, "old");
- server.updateAllOutput(new OutputUpdateAllEvent("note", "para", Collections.singletonList(
- new RemoteInterpreterResultMessage("HTML", "replacement")), null));
- verify(listener).onParagraphOutputUpdated(
- "note", "para", 0, null, InterpreterResult.Type.HTML, "replacement");
- runner.appendBuffer("note", "para", 0, null, "new");
- runner.run();
+ server.appendOutput(new OutputAppendEvent("note", "para", 0, "old", null, null));
+ server.updateAllOutput(new OutputUpdateAllEvent("note", "para", Arrays.asList(
+ new RemoteInterpreterResultMessage("HTML", "replacement"),
+ new RemoteInterpreterResultMessage("TEXT", "second")), null));
+ verify(listener).onParagraphOutputUpdated("note", "para", 1, null, InterpreterResult.Type.TEXT, "second");
+ server.appendOutput(new OutputAppendEvent("note", "para", 1, "new", null, null));
+ server.checkpointOutput("note", "para");
InOrder order = inOrder(listener);
order.verify(listener).onParagraphOutputAppend("note", "para", 0, null, "old");
order.verify(listener).onParagraphOutputClear("note", "para", null);
- order.verify(listener).onParagraphOutputUpdated(
- "note", "para", 0, null, InterpreterResult.Type.HTML, "replacement");
- order.verify(listener).onParagraphOutputAppend("note", "para", 0, null, "new");
+ order.verify(listener).onParagraphOutputUpdated("note", "para", 0, null, InterpreterResult.Type.HTML, "replacement");
+ order.verify(listener).onParagraphOutputUpdated("note", "para", 1, null, InterpreterResult.Type.TEXT, "second");
+ order.verify(listener).onParagraphOutputAppend("note", "para", 1, null, "new");
+ order.verify(listener).checkpointOutput("note", "para");
+ order.verifyNoMoreInteractions();
} finally {
server.stop();
}
}
@Test
- void updateAllWaitsForInFlightAppendAndCompletesBeforeReturning() throws Exception {
+ void emptyUpdateAllStillClearsPendingOutput() throws Exception {
+ RemoteInterpreterProcessListener listener = mock(RemoteInterpreterProcessListener.class);
+ RemoteInterpreterEventServer server = serverWithListener(listener);
+ try {
+ server.appendOutput(new OutputAppendEvent("note", "para", 0, "old", null, null));
+ server.updateAllOutput(new OutputUpdateAllEvent("note", "para", Collections.emptyList(), null));
+ InOrder order = inOrder(listener);
+ order.verify(listener).onParagraphOutputAppend("note", "para", 0, null, "old");
+ order.verify(listener).onParagraphOutputClear("note", "para", null);
+ order.verifyNoMoreInteractions();
+ } finally {
+ server.stop();
+ }
+ }
+
+ @ParameterizedTest
+ @ValueSource(strings = {"UPDATE", "UPDATE_ALL", "CHECKPOINT"})
+ void sameNoteBoundaryWaitsForInFlightAppend(String operation) throws Exception {
RemoteInterpreterProcessListener listener = mock(RemoteInterpreterProcessListener.class);
- AppendOutputRunner runner = new AppendOutputRunner(listener);
- RemoteInterpreterEventServer server = serverWithRunner(listener, runner);
+ RemoteInterpreterEventServer server = serverWithListener(listener);
CountDownLatch entered = new CountDownLatch(1);
CountDownLatch release = new CountDownLatch(1);
- CountDownLatch updateStarted = new CountDownLatch(1);
+ CountDownLatch started = new CountDownLatch(1);
doAnswer(invocation -> {
entered.countDown();
assertTrue(release.await(5, TimeUnit.SECONDS));
return null;
- }).when(listener).onParagraphOutputAppend("note", "para", 0, null, "old");
- ExecutorService executor = Executors.newFixedThreadPool(2);
+ }).when(listener).onParagraphOutputAppend("note", "first", 0, null, "old");
+ ExecutorService callers = Executors.newSingleThreadExecutor();
try {
- runner.appendBuffer("note", "para", 0, null, "old");
- Future first = executor.submit(runner);
+ server.appendOutput(new OutputAppendEvent("note", "first", 0, "old", null, null));
+ dispatcherOf(server).flush();
assertTrue(entered.await(5, TimeUnit.SECONDS));
- Future update = executor.submit(() -> {
- updateStarted.countDown();
- server.updateAllOutput(new OutputUpdateAllEvent("note", "para", Collections.singletonList(
- new RemoteInterpreterResultMessage("HTML", "replacement")), null));
- verify(listener).onParagraphOutputUpdated(
- "note", "para", 0, null, InterpreterResult.Type.HTML, "replacement");
+ Future boundary = callers.submit(() -> {
+ started.countDown();
+ callBoundary(server, "note", "second", operation);
return null;
});
- assertTrue(updateStarted.await(5, TimeUnit.SECONDS));
- assertThrows(TimeoutException.class, () -> update.get(100, TimeUnit.MILLISECONDS));
+ assertTrue(started.await(5, TimeUnit.SECONDS));
+ assertThrows(TimeoutException.class, () -> boundary.get(100, TimeUnit.MILLISECONDS));
release.countDown();
- first.get(5, TimeUnit.SECONDS);
- update.get(5, TimeUnit.SECONDS);
+ boundary.get(5, TimeUnit.SECONDS);
InOrder order = inOrder(listener);
- order.verify(listener).onParagraphOutputAppend("note", "para", 0, null, "old");
+ order.verify(listener).onParagraphOutputAppend("note", "first", 0, null, "old");
+ if ("UPDATE".equals(operation)) {
+ order.verify(listener).onParagraphOutputUpdated("note", "second", 0, null, InterpreterResult.Type.TEXT, "replacement");
+ } else if ("UPDATE_ALL".equals(operation)) {
+ order.verify(listener).onParagraphOutputClear("note", "second", null);
+ order.verify(listener).onParagraphOutputUpdated("note", "second", 0, null, InterpreterResult.Type.TEXT, "replacement");
+ } else {
+ order.verify(listener).checkpointOutput("note", "second");
+ }
+ } finally {
+ release.countDown();
+ callers.shutdownNow();
+ server.stop();
+ }
+ }
+
+ @Test
+ void independentNoteMakesProgressWhileAnotherNoteAppendIsBlocked() throws Exception {
+ // These note IDs have identical hash codes, so concurrency cannot depend on hash lanes.
+ RemoteInterpreterProcessListener listener = mock(RemoteInterpreterProcessListener.class);
+ RemoteInterpreterEventServer server = serverWithListener(listener);
+ CountDownLatch entered = new CountDownLatch(1);
+ CountDownLatch release = new CountDownLatch(1);
+ CountDownLatch otherAppend = new CountDownLatch(1);
+ doAnswer(call -> {
+ entered.countDown();
+ assertTrue(release.await(5, TimeUnit.SECONDS));
+ return null;
+ }).when(listener).onParagraphOutputAppend("Aa", "para", 0, null, "blocked");
+ doAnswer(call -> {
+ otherAppend.countDown();
+ return null;
+ }).when(listener).onParagraphOutputAppend("BB", "para", 0, null, "periodic");
+ ExecutorService callers = Executors.newSingleThreadExecutor();
+ try {
+ server.appendOutput(new OutputAppendEvent("Aa", "para", 0, "blocked", null, null));
+ dispatcherOf(server).flush();
+ assertTrue(entered.await(5, TimeUnit.SECONDS));
+ server.appendOutput(new OutputAppendEvent("BB", "para", 0, "periodic", null, null));
+ dispatcherOf(server).flush();
+ assertTrue(otherAppend.await(5, TimeUnit.SECONDS));
+ Future independent = callers.submit(() -> {
+ callBoundary(server, "BB", "para", "UPDATE");
+ callBoundary(server, "BB", "para", "UPDATE_ALL");
+ callBoundary(server, "BB", "para", "CHECKPOINT");
+ return null;
+ });
+ independent.get(5, TimeUnit.SECONDS);
+ verify(listener).onParagraphOutputClear("BB", "para", null);
+ verify(listener).checkpointOutput("BB", "para");
+ } finally {
+ release.countDown();
+ callers.shutdownNow();
+ server.stop();
+ }
+ }
+
+ @ParameterizedTest
+ @ValueSource(strings = {"UPDATE", "UPDATE_ALL", "CHECKPOINT"})
+ void laterAppendCannotOvertakeBoundaryCallback(String operation) throws Exception {
+ RemoteInterpreterProcessListener listener = mock(RemoteInterpreterProcessListener.class);
+ RemoteInterpreterEventServer server = serverWithListener(listener);
+ CountDownLatch entered = new CountDownLatch(1);
+ CountDownLatch release = new CountDownLatch(1);
+ AtomicBoolean boundaryCompleted = new AtomicBoolean();
+ AtomicBoolean appendOverlappedBoundary = new AtomicBoolean();
+ doAnswer(call -> {
+ entered.countDown();
+ assertTrue(release.await(5, TimeUnit.SECONDS));
+ return null;
+ }).when(listener).onParagraphOutputClear("note", "para", null);
+ doAnswer(call -> {
+ entered.countDown();
+ assertTrue(release.await(5, TimeUnit.SECONDS));
+ boundaryCompleted.set(true);
+ return null;
+ }).when(listener).checkpointOutput("note", "para");
+ doAnswer(call -> {
+ entered.countDown();
+ assertTrue(release.await(5, TimeUnit.SECONDS));
+ boundaryCompleted.set(true);
+ return null;
+ }).when(listener).onParagraphOutputUpdated("note", "para", 0, null, InterpreterResult.Type.TEXT, "replacement");
+ doAnswer(call -> {
+ if (!boundaryCompleted.get()) {
+ appendOverlappedBoundary.set(true);
+ }
+ return null;
+ }).when(listener).onParagraphOutputAppend("note", "para", 0, null, "later");
+ ExecutorService callers = Executors.newSingleThreadExecutor();
+ try {
+ Future boundary = callers.submit(() -> {
+ callBoundary(server, "note", "para", operation);
+ return null;
+ });
+ assertTrue(entered.await(5, TimeUnit.SECONDS));
+ server.appendOutput(new OutputAppendEvent("note", "para", 0, "later", null, null));
+ dispatcherOf(server).flush();
+ verify(listener, never()).onParagraphOutputAppend("note", "para", 0, null, "later");
+ assertThrows(TimeoutException.class, () -> boundary.get(100, TimeUnit.MILLISECONDS));
+ release.countDown();
+ boundary.get(5, TimeUnit.SECONDS);
+ server.checkpointOutput("note", "drained");
+ assertTrue(boundaryCompleted.get());
+ assertFalse(appendOverlappedBoundary.get(), "Append must wait for the entire boundary");
+ InOrder order = inOrder(listener);
+ if ("UPDATE".equals(operation)) {
+ order.verify(listener).onParagraphOutputUpdated("note", "para", 0, null, InterpreterResult.Type.TEXT, "replacement");
+ } else if ("UPDATE_ALL".equals(operation)) {
+ order.verify(listener).onParagraphOutputClear("note", "para", null);
+ order.verify(listener).onParagraphOutputUpdated("note", "para", 0, null, InterpreterResult.Type.TEXT, "replacement");
+ } else {
+ order.verify(listener).checkpointOutput("note", "para");
+ }
+ order.verify(listener).onParagraphOutputAppend("note", "para", 0, null, "later");
+ } finally {
+ release.countDown();
+ callers.shutdownNow();
+ server.stop();
+ }
+ }
+
+ @ParameterizedTest
+ @ValueSource(strings = {"UPDATE", "UPDATE_ALL", "CHECKPOINT"})
+ void boundaryCallbackFailureBecomesRpcExceptionAndLaterOutputStillWorks(String operation)
+ throws Exception {
+ RemoteInterpreterProcessListener listener = mock(RemoteInterpreterProcessListener.class);
+ IllegalStateException removed = new IllegalStateException("paragraph removed");
+ doThrow(removed).when(listener).onParagraphOutputUpdated("note", "gone", 0, null, InterpreterResult.Type.TEXT, "replacement");
+ doThrow(removed).when(listener).onParagraphOutputClear("note", "gone", null);
+ doThrow(removed).when(listener).checkpointOutput("note", "gone");
+ RemoteInterpreterEventServer server = serverWithListener(listener);
+ try {
+ InterpreterRPCException failure = assertThrows(InterpreterRPCException.class,
+ () -> callBoundary(server, "note", "gone", operation));
+ assertTrue(failure.getErrorMessage().contains("paragraph removed"));
+ server.appendOutput(new OutputAppendEvent("note", "present", 0, "good", null, null));
+ server.checkpointOutput("note", "present");
+ verify(listener).onParagraphOutputAppend("note", "present", 0, null, "good");
+ } finally {
+ server.stop();
+ }
+ }
+
+ @Test
+ void largeParagraphOutputIsDeliveredWithoutAnAdditionalDispatcherLimit() throws Exception {
+ RemoteInterpreterProcessListener listener = mock(RemoteInterpreterProcessListener.class);
+ RemoteInterpreterEventServer server = serverWithListener(listener);
+ String output = "a".repeat(4 * 1024 * 1024 + 1);
+ try {
+ server.appendOutput(new OutputAppendEvent("note", "para", 0, output, null, null));
+ server.updateOutput(new OutputUpdateEvent("note", "para", 0, "TEXT", output, null, null));
+ server.updateAllOutput(new OutputUpdateAllEvent("note", "para", List.of(
+ new RemoteInterpreterResultMessage("HTML", output)), null));
+ server.checkpointOutput("note", "para");
+ InOrder order = inOrder(listener);
+ order.verify(listener).onParagraphOutputAppend("note", "para", 0, null, output);
+ order.verify(listener).onParagraphOutputUpdated("note", "para", 0, null, InterpreterResult.Type.TEXT, output);
order.verify(listener).onParagraphOutputClear("note", "para", null);
- order.verify(listener).onParagraphOutputUpdated(
- "note", "para", 0, null, InterpreterResult.Type.HTML, "replacement");
+ order.verify(listener).onParagraphOutputUpdated("note", "para", 0, null, InterpreterResult.Type.HTML, output);
+ order.verify(listener).checkpointOutput("note", "para");
+ order.verifyNoMoreInteractions();
} finally {
+ server.stop();
+ }
+ }
+
+ @Test
+ void serverAppliesConfiguredOutputBatchSize() throws Exception {
+ ZeppelinConfiguration zConf = ZeppelinConfiguration.load("zeppelin-test-site.xml");
+ zConf.setProperty(ConfVars.ZEPPELIN_INTERPRETER_OUTPUT_WORKER_COUNT.getVarName(), "1");
+ zConf.setProperty(ConfVars.ZEPPELIN_INTERPRETER_OUTPUT_EVENTS_PER_BATCH.getVarName(), "2");
+ RemoteInterpreterProcessListener listener = mock(RemoteInterpreterProcessListener.class);
+ RemoteInterpreterEventServer server = serverWithListener(listener, zConf);
+ try {
+ for (String output : List.of("a", "b", "c", "d")) {
+ server.appendOutput(new OutputAppendEvent("note", "para", 0, output, null, null));
+ }
+ server.checkpointOutput("note", "para");
+ InOrder order = inOrder(listener);
+ order.verify(listener).onParagraphOutputAppend("note", "para", 0, null, "ab");
+ order.verify(listener).onParagraphOutputAppend("note", "para", 0, null, "cd");
+ order.verify(listener).checkpointOutput("note", "para");
+ order.verifyNoMoreInteractions();
+ } finally {
+ server.stop();
+ }
+ }
+
+ @Test
+ void stoppedServerRejectsOutputWithRpcExceptions() {
+ RemoteInterpreterProcessListener listener = mock(RemoteInterpreterProcessListener.class);
+ RemoteInterpreterEventServer server = serverWithListener(listener);
+ server.stop();
+ assertThrows(InterpreterRPCException.class, () ->
+ server.appendOutput(new OutputAppendEvent("note", "para", 0, "late", null, null)));
+ for (String operation : Arrays.asList("UPDATE", "UPDATE_ALL", "CHECKPOINT")) {
+ assertThrows(InterpreterRPCException.class,
+ () -> callBoundary(server, "note", "para", operation));
+ }
+ verify(listener, never()).onParagraphOutputAppend("note", "para", 0, null, "late");
+ }
+
+ @Test
+ void interruptedRpcWaitPreservesInterruptAndDoesNotCancelAcceptedOutput() throws Exception {
+ RemoteInterpreterProcessListener listener = mock(RemoteInterpreterProcessListener.class);
+ CountDownLatch entered = new CountDownLatch(1);
+ CountDownLatch release = new CountDownLatch(1);
+ doAnswer(call -> {
+ entered.countDown();
+ assertTrue(release.await(5, TimeUnit.SECONDS));
+ return null;
+ }).when(listener).checkpointOutput("note", "para");
+ RemoteInterpreterEventServer server = serverWithListener(listener);
+ AtomicBoolean interruptPreserved = new AtomicBoolean();
+ AtomicReference failure = new AtomicReference<>();
+ Thread caller = new Thread(() -> {
+ try {
+ server.checkpointOutput("note", "para");
+ } catch (Throwable e) {
+ failure.set(e);
+ interruptPreserved.set(Thread.currentThread().isInterrupted());
+ }
+ });
+ try {
+ caller.start();
+ assertTrue(entered.await(5, TimeUnit.SECONDS));
+ caller.interrupt();
+ caller.join(5000);
+ assertFalse(caller.isAlive());
+ assertTrue(failure.get() instanceof InterpreterRPCException);
+ assertTrue(interruptPreserved.get());
release.countDown();
- executor.shutdownNow();
+ server.appendOutput(new OutputAppendEvent("note", "para", 0, "after", null, null));
+ server.checkpointOutput("note", "done");
+ InOrder order = inOrder(listener);
+ order.verify(listener).checkpointOutput("note", "para");
+ order.verify(listener).onParagraphOutputAppend("note", "para", 0, null, "after");
+ order.verify(listener).checkpointOutput("note", "done");
+ } finally {
+ release.countDown();
+ caller.interrupt();
+ caller.join(5000);
server.stop();
}
}
- private RemoteInterpreterEventServer serverWithRunner(
- RemoteInterpreterProcessListener listener, AppendOutputRunner runner) throws Exception {
+ private void callBoundary(RemoteInterpreterEventServer server, String noteId,
+ String paragraphId, String operation) throws Exception {
+ if ("UPDATE".equals(operation)) {
+ server.updateOutput(new OutputUpdateEvent(noteId, paragraphId, 0, "TEXT", "replacement", null, null));
+ } else if ("UPDATE_ALL".equals(operation)) {
+ server.updateAllOutput(new OutputUpdateAllEvent(noteId, paragraphId, Collections.singletonList(new RemoteInterpreterResultMessage("TEXT", "replacement")), null));
+ } else {
+ server.checkpointOutput(noteId, paragraphId);
+ }
+ }
+
+ private RemoteInterpreterEventServer serverWithListener(
+ RemoteInterpreterProcessListener listener) {
+ return serverWithListener(listener, outputConfiguration());
+ }
+
+ private RemoteInterpreterEventServer serverWithListener(
+ RemoteInterpreterProcessListener listener, ZeppelinConfiguration zConf) {
InterpreterSettingManager manager = mock(InterpreterSettingManager.class);
when(manager.getRemoteInterpreterProcessListener()).thenReturn(listener);
- RemoteInterpreterEventServer server = new RemoteInterpreterEventServer(
- mock(ZeppelinConfiguration.class), manager);
- Field field = RemoteInterpreterEventServer.class.getDeclaredField("runner");
+ return new RemoteInterpreterEventServer(zConf, manager);
+ }
+
+ private ZeppelinConfiguration outputConfiguration() {
+ ZeppelinConfiguration zConf = mock(ZeppelinConfiguration.class);
+ when(zConf.getInt(ConfVars.ZEPPELIN_INTERPRETER_OUTPUT_WORKER_COUNT))
+ .thenReturn(ConfVars.ZEPPELIN_INTERPRETER_OUTPUT_WORKER_COUNT.getIntValue());
+ when(zConf.getInt(ConfVars.ZEPPELIN_INTERPRETER_OUTPUT_EVENTS_PER_BATCH))
+ .thenReturn(ConfVars.ZEPPELIN_INTERPRETER_OUTPUT_EVENTS_PER_BATCH.getIntValue());
+ return zConf;
+ }
+
+ private ParagraphOutputDispatcher dispatcherOf(RemoteInterpreterEventServer server)
+ throws Exception {
+ Field field = RemoteInterpreterEventServer.class.getDeclaredField("outputDispatcher");
field.setAccessible(true);
- field.set(server, runner);
- return server;
+ return (ParagraphOutputDispatcher) field.get(server);
}
@Test
void invokeMethodThrowsRpcExceptionWhenSerializationFails() throws Exception {
- ZeppelinConfiguration zConf = mock(ZeppelinConfiguration.class);
+ ZeppelinConfiguration zConf = outputConfiguration();
InterpreterSettingManager manager = mock(InterpreterSettingManager.class);
RemoteInterpreterEventServer server = new RemoteInterpreterEventServer(zConf, manager);
@@ -189,7 +503,8 @@ void invokeMethodThrowsRpcExceptionWhenSerializationFails() throws Exception {
.callRemoteFunction(any());
ResourceId resourceId = ResourceId.fromJson(
- "{\"resourcePoolId\":\"pool-id\",\"name\":\"resource-name\",\"noteId\":\"note-id\",\"paragraphId\":\"paragraph-id\"}"
+ "{\"resourcePoolId\":\"pool-id\",\"name\":\"resource-name\","
+ + "\"noteId\":\"note-id\",\"paragraphId\":\"paragraph-id\"}"
);
InvokeResourceMethodEventMessage message = new InvokeResourceMethodEventMessage(
diff --git a/zeppelin-server/src/test/java/org/apache/zeppelin/interpreter/remote/AppendOutputRunnerTest.java b/zeppelin-server/src/test/java/org/apache/zeppelin/interpreter/remote/AppendOutputRunnerTest.java
index 786c70e23c8..7f51d080d35 100644
--- a/zeppelin-server/src/test/java/org/apache/zeppelin/interpreter/remote/AppendOutputRunnerTest.java
+++ b/zeppelin-server/src/test/java/org/apache/zeppelin/interpreter/remote/AppendOutputRunnerTest.java
@@ -17,366 +17,182 @@
package org.apache.zeppelin.interpreter.remote;
-import org.apache.zeppelin.interpreter.InterpreterResult;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.ArgumentMatchers.isNull;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.inOrder;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoInteractions;
import org.apache.log4j.AppenderSkeleton;
import org.apache.log4j.Level;
import org.apache.log4j.Logger;
import org.apache.log4j.spi.LoggingEvent;
-import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
import org.mockito.InOrder;
-import org.mockito.invocation.InvocationOnMock;
-import org.mockito.stubbing.Answer;
-
import java.util.ArrayList;
import java.util.List;
-import java.time.Duration;
-import java.util.concurrent.CountDownLatch;
-import java.util.concurrent.Future;
-import java.util.concurrent.ExecutorService;
-import java.util.concurrent.TimeoutException;
-import java.util.concurrent.Executors;
-import java.util.concurrent.ScheduledExecutorService;
-import java.util.concurrent.ScheduledFuture;
-import java.util.concurrent.TimeUnit;
-
-import static org.junit.jupiter.api.Assertions.assertTimeoutPreemptively;
-import static org.junit.jupiter.api.Assertions.assertTrue;
-import static org.junit.jupiter.api.Assertions.assertThrows;
-import static org.mockito.Mockito.doThrow;
-import static org.junit.jupiter.api.Assertions.fail;
-import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.ArgumentMatchers.anyInt;
-import static org.mockito.ArgumentMatchers.isNull;
-import static org.mockito.Mockito.atMost;
-import static org.mockito.Mockito.doAnswer;
-import static org.mockito.Mockito.inOrder;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.times;
-import static org.mockito.Mockito.verify;
class AppendOutputRunnerTest {
-
- private static final int NUM_EVENTS = 10000;
- private static final int NUM_CLUBBED_EVENTS = 100;
- private static final ScheduledExecutorService service = Executors.newSingleThreadScheduledExecutor();
- private static ScheduledFuture future = null;
- /* It is being accessed by multiple threads.
- * While loop for 'loopForBufferCompletion' could
- * run for-ever.
- */
- private volatile static int numInvocations = 0;
-
- @AfterEach
- public void afterEach() {
- if (future != null) {
- future.cancel(true);
- }
- }
-
- @Test
- void testSingleEvent() throws InterruptedException {
- RemoteInterpreterProcessListener listener = mock(RemoteInterpreterProcessListener.class);
- String[][] buffer = {{"note", "para", "data\n"}};
-
- loopForCompletingEvents(listener, 1, buffer);
- verify(listener, times(1)).onParagraphOutputAppend(
- any(String.class), any(String.class), anyInt(), isNull(), any(String.class));
- verify(listener, times(1)).onParagraphOutputAppend("note", "para", 0, null, "data\n");
- }
-
@Test
- public void testMultipleEventsOfSameParagraph() throws InterruptedException {
- RemoteInterpreterProcessListener listener = mock(RemoteInterpreterProcessListener.class);
- String note1 = "note1";
- String para1 = "para1";
- String[][] buffer = {
- {note1, para1, "data1\n"},
- {note1, para1, "data2\n"},
- {note1, para1, "data3\n"}
- };
-
- loopForCompletingEvents(listener, 1, buffer);
- verify(listener, times(1)).onParagraphOutputAppend(
- any(String.class), any(String.class), anyInt(), isNull(), any(String.class));
- verify(listener, times(1)).onParagraphOutputAppend(
- note1, para1, 0, null, "data1\ndata2\ndata3\n");
- }
-
- // A paragraph in shared mode is one Job, so one execution -- one user -- produces the output
- // for a given index at a time, and keying by user must leave that batching alone.
- @Test
- void appendsFromOneExecutionAreStillBatchedIntoOneChunk() {
+ void batchesAdjacentAppendsAndHandlesEmptyBatches() {
RemoteInterpreterProcessListener listener = mock(RemoteInterpreterProcessListener.class);
AppendOutputRunner runner = new AppendOutputRunner(listener);
- runner.appendBuffer("note", "para", 0, "user1", "line-1\n");
- runner.appendBuffer("note", "para", 0, "user1", "line-2\n");
-
- runner.run();
-
- verify(listener, times(1)).onParagraphOutputAppend(
- "note", "para", 0, "user1", "line-1\nline-2\n");
+ List batch = new ArrayList<>();
+ batch.add(new AppendOutputBuffer("note", "para", 0, null, "a"));
+ batch.add(new AppendOutputBuffer("note", "para", 0, null, "b"));
+ verifyNoInteractions(listener);
+ runner.run(batch);
+ batch.clear();
+ batch.add(new AppendOutputBuffer("note", "para", 0, null, "c"));
+ runner.run(batch);
+ batch.clear();
+ runner.run(batch);
+ InOrder order = inOrder(listener);
+ order.verify(listener).onParagraphOutputAppend("note", "para", 0, null, "ab");
+ order.verify(listener).onParagraphOutputAppend("note", "para", 0, null, "c");
+ order.verifyNoMoreInteractions();
}
- // Two executions only overlap in personalized mode, where each user runs their own copy.
- // A merged chunk would have no single owner and could not be routed to either user.
@Test
- void appendsFromDifferentExecutionsAreNotBatchedTogether() {
+ void mergesEachOutputKeyInFirstAppearanceOrder() {
RemoteInterpreterProcessListener listener = mock(RemoteInterpreterProcessListener.class);
AppendOutputRunner runner = new AppendOutputRunner(listener);
- runner.appendBuffer("note", "para", 0, "user1", "mine\n");
- runner.appendBuffer("note", "para", 0, "user2", "theirs\n");
-
- runner.run();
-
- verify(listener, times(1)).onParagraphOutputAppend("note", "para", 0, "user1", "mine\n");
- verify(listener, times(1)).onParagraphOutputAppend("note", "para", 0, "user2", "theirs\n");
+ List batch = new ArrayList<>();
+ batch.add(new AppendOutputBuffer("note:1", "p:1", 0, null, "first"));
+ batch.add(new AppendOutputBuffer("note:1", "p:2", 0, null, "second"));
+ batch.add(new AppendOutputBuffer("note:1", "p:1", 1, null, "third"));
+ batch.add(new AppendOutputBuffer("note:2", "p:1", 1, null, "fourth"));
+ batch.add(new AppendOutputBuffer("note:1", "p:1", 0, null, "fifth"));
+ runner.run(batch);
+ InOrder order = inOrder(listener);
+ order.verify(listener).onParagraphOutputAppend("note:1", "p:1", 0, null, "firstfifth");
+ order.verify(listener).onParagraphOutputAppend("note:1", "p:2", 0, null, "second");
+ order.verify(listener).onParagraphOutputAppend("note:1", "p:1", 1, null, "third");
+ order.verify(listener).onParagraphOutputAppend("note:2", "p:1", 1, null, "fourth");
+ order.verifyNoMoreInteractions();
}
- // Keying by user splits the buffer, so the per-owner ordering the shared queue guarantees
- // must still hold once another user's output is interleaved.
@Test
- void updatesDoNotOvertakeQueuedAppendsOfTheSameOwner() {
+ void doesNotMergeOutputFromDifferentExecutionOwners() {
RemoteInterpreterProcessListener listener = mock(RemoteInterpreterProcessListener.class);
AppendOutputRunner runner = new AppendOutputRunner(listener);
- runner.appendBuffer("note", "para", 0, "owner", "before\n");
- runner.appendBuffer("note", "para", 0, "other", "theirs\n");
- runner.updateBuffer("note", "para", 0, "owner", InterpreterResult.Type.TEXT, "replacement\n");
- runner.appendBuffer("note", "para", 0, "owner", "after\n");
+ List batch = new ArrayList<>();
+ batch.add(new AppendOutputBuffer("note", "para", 0, "alice", "a"));
+ batch.add(new AppendOutputBuffer("note", "para", 0, "bob", "b"));
+ batch.add(new AppendOutputBuffer("note", "para", 0, "alice", "c"));
- runner.run();
+ runner.run(batch);
InOrder order = inOrder(listener);
- order.verify(listener).onParagraphOutputAppend("note", "para", 0, "owner", "before\n");
- order.verify(listener).onParagraphOutputUpdated(
- "note", "para", 0, "owner", InterpreterResult.Type.TEXT, "replacement\n");
- order.verify(listener).onParagraphOutputAppend("note", "para", 0, "owner", "after\n");
+ order.verify(listener).onParagraphOutputAppend("note", "para", 0, "alice", "ac");
+ order.verify(listener).onParagraphOutputAppend("note", "para", 0, "bob", "b");
+ order.verifyNoMoreInteractions();
}
@Test
- void testUpdateDoesNotOvertakeQueuedAppend() {
+ void staleAppendDoesNotDiscardLaterOutputOrLeakIntoItsBuffer() {
RemoteInterpreterProcessListener listener = mock(RemoteInterpreterProcessListener.class);
+ doThrow(new IllegalStateException("removed")).when(listener)
+ .onParagraphOutputAppend("note", "gone", 0, null, "bad");
AppendOutputRunner runner = new AppendOutputRunner(listener);
- runner.appendBuffer("note", "para", 0, null, "before-1\n");
- runner.appendBuffer("note", "para", 0, null, "before-2\n");
- runner.updateBuffer("note", "para", 0, null, InterpreterResult.Type.TEXT, "replacement\n");
- runner.appendBuffer("note", "para", 0, null, "after\n");
-
- runner.run();
-
+ List batch = new ArrayList<>();
+ batch.add(new AppendOutputBuffer("note", "gone", 0, null, "bad"));
+ batch.add(new AppendOutputBuffer("note", "present", 0, null, "good"));
+ runner.run(batch);
InOrder order = inOrder(listener);
- order.verify(listener).onParagraphOutputAppend("note", "para", 0, null, "before-1\nbefore-2\n");
- order.verify(listener).onParagraphOutputUpdated(
- "note", "para", 0, null, InterpreterResult.Type.TEXT, "replacement\n");
- order.verify(listener).onParagraphOutputAppend("note", "para", 0, null, "after\n");
- }
-
- @Test
- void testMultipleEventsOfDifferentParagraphs() throws InterruptedException {
- RemoteInterpreterProcessListener listener = mock(RemoteInterpreterProcessListener.class);
- String note1 = "note1";
- String note2 = "note2";
- String para1 = "para1";
- String para2 = "para2";
- String[][] buffer = {
- {note1, para1, "data1\n"},
- {note1, para2, "data2\n"},
- {note2, para1, "data3\n"},
- {note2, para2, "data4\n"}
- };
- loopForCompletingEvents(listener, 4, buffer);
-
- verify(listener, times(4)).onParagraphOutputAppend(
- any(String.class), any(String.class), anyInt(), isNull(), any(String.class));
- verify(listener, times(1)).onParagraphOutputAppend(note1, para1, 0, null, "data1\n");
- verify(listener, times(1)).onParagraphOutputAppend(note1, para2, 0, null, "data2\n");
- verify(listener, times(1)).onParagraphOutputAppend(note2, para1, 0, null, "data3\n");
- verify(listener, times(1)).onParagraphOutputAppend(note2, para2, 0, null, "data4\n");
+ order.verify(listener).onParagraphOutputAppend("note", "gone", 0, null, "bad");
+ order.verify(listener).onParagraphOutputAppend("note", "present", 0, null, "good");
+ order.verifyNoMoreInteractions();
}
@Test
- void testClubbedData() throws InterruptedException {
+ void largeAppendStreamIsDeliveredAsOneBatchWithoutLosingData() {
RemoteInterpreterProcessListener listener = mock(RemoteInterpreterProcessListener.class);
+ List received = new ArrayList<>();
+ doAnswer(call -> {
+ received.add(call.getArgument(4));
+ return null;
+ }).when(listener).onParagraphOutputAppend(
+ anyString(), anyString(), anyInt(), isNull(), anyString());
AppendOutputRunner runner = new AppendOutputRunner(listener);
- future = service.scheduleWithFixedDelay(runner, 0,
- AppendOutputRunner.BUFFER_TIME_MS, TimeUnit.MILLISECONDS);
- Thread thread = new Thread(new BombardEvents(runner));
- thread.start();
- thread.join();
- Thread.sleep(1000);
-
- /* NUM_CLUBBED_EVENTS is a heuristic number.
- * It has been observed that for 10,000 continuos event
- * calls, 30-40 Web-socket calls are made. Keeping
- * the unit-test to a pessimistic 100 web-socket calls.
- */
- verify(listener, atMost(NUM_CLUBBED_EVENTS)).onParagraphOutputAppend(
- any(String.class), any(String.class), anyInt(), isNull(), any(String.class));
+ List batch = new ArrayList<>();
+ StringBuilder expected = new StringBuilder();
+ for (int i = 0; i < 10000; i++) {
+ String token = i + "\n";
+ expected.append(token);
+ batch.add(new AppendOutputBuffer("note", "para", 0, null, token));
+ }
+ runner.run(batch);
+ assertEquals(List.of(expected.toString()), received);
}
- @Test
- void testWarnLoggerForLargeData() throws InterruptedException {
+ @ParameterizedTest
+ @ValueSource(ints = {100000, 100001})
+ void warnsOnlyWhenBufferedOutputExceedsTheSizeThreshold(int size) {
RemoteInterpreterProcessListener listener = mock(RemoteInterpreterProcessListener.class);
AppendOutputRunner runner = new AppendOutputRunner(listener);
- String data = "data\n";
- int numEvents = 100000;
+ List batch = new ArrayList<>();
+ String output = "a".repeat(size);
+ batch.add(new AppendOutputBuffer("note", "para", 0, null, output));
+ List sizeWarnings = new ArrayList<>();
+ AppenderSkeleton appender = new AppenderSkeleton() {
+ @Override
+ protected void append(LoggingEvent event) {
+ String message = event.getRenderedMessage();
+ if (Level.WARN.equals(event.getLevel())
+ && message.startsWith("Processing size for buffered append-output is high:")) {
+ sizeWarnings.add(message);
+ }
+ }
- for (int i=0; i
- Level.WARN.equals(event.getLevel()) && expected.equals(event.getMessage())));
+ runner.run(batch);
+ assertEquals(size > 100000 ? List.of(
+ "Processing size for buffered append-output is high: " + size + " characters.")
+ : List.of(), sizeWarnings);
+ verify(listener).onParagraphOutputAppend("note", "para", 0, null, output);
} finally {
logger.removeAppender(appender);
+ logger.setLevel(previousLevel);
+ logger.setAdditivity(previousAdditivity);
+ appender.close();
}
}
@Test
- void emptyDrainDoesNotBlock() {
- AppendOutputRunner runner =
- new AppendOutputRunner(mock(RemoteInterpreterProcessListener.class));
- assertTimeoutPreemptively(Duration.ofSeconds(1), runner::run);
- }
-
- @Test
- void updateFailureDoesNotDiscardOtherEventsOrLaterDrains() {
- RemoteInterpreterProcessListener listener = mock(RemoteInterpreterProcessListener.class);
- AppendOutputRunner runner = new AppendOutputRunner(listener);
- doThrow(new IllegalStateException("removed")).when(listener)
- .onParagraphOutputUpdated("note", "gone", 0, null, InterpreterResult.Type.TEXT, "bad");
- runner.appendBuffer("note", "gone", 0, null, "bad");
- runner.updateBuffer("note", "gone", 0, null, InterpreterResult.Type.TEXT, "bad");
- runner.appendBuffer("note", "present", 0, null, "good");
- runner.run();
- runner.appendBuffer("note", "present", 0, null, "later");
- runner.run();
- verify(listener).onParagraphOutputAppend("note", "present", 0, null, "good");
- verify(listener).onParagraphOutputAppend("note", "present", 0, null, "later");
- }
-
- @Test
- void appendFailureDoesNotDiscardLaterUpdateOrDrain() {
+ void disallowedDeliverySkipsTheBatchAndDoesNotAffectLaterBatches() {
RemoteInterpreterProcessListener listener = mock(RemoteInterpreterProcessListener.class);
AppendOutputRunner runner = new AppendOutputRunner(listener);
- doThrow(new IllegalStateException("removed")).when(listener)
- .onParagraphOutputAppend("note", "gone", 0, null, "bad");
- runner.appendBuffer("note", "gone", 0, null, "bad");
- runner.updateBuffer("note", "present", 0, null, InterpreterResult.Type.TEXT, "current");
- runner.run();
- runner.appendBuffer("note", "present", 0, null, "later");
- runner.run();
- verify(listener).onParagraphOutputUpdated(
- "note", "present", 0, null, InterpreterResult.Type.TEXT, "current");
- verify(listener).onParagraphOutputAppend("note", "present", 0, null, "later");
- }
-
- @Test
- void concurrentDrainCannotOvertakeInFlightCallback() throws Exception {
- RemoteInterpreterProcessListener listener = mock(RemoteInterpreterProcessListener.class);
- AppendOutputRunner runner = new AppendOutputRunner(listener);
- CountDownLatch entered = new CountDownLatch(1);
- CountDownLatch release = new CountDownLatch(1);
- doAnswer(invocation -> {
- entered.countDown();
- assertTrue(release.await(5, TimeUnit.SECONDS));
- return null;
- }).when(listener).onParagraphOutputAppend("note", "para", 0, null, "old");
- ExecutorService executor = Executors.newFixedThreadPool(2);
- try {
- runner.appendBuffer("note", "para", 0, null, "old");
- Future first = executor.submit(runner);
- assertTrue(entered.await(5, TimeUnit.SECONDS));
- runner.updateBuffer("note", "para", 0, null, InterpreterResult.Type.TEXT, "new");
- Future second = executor.submit(runner);
- assertThrows(TimeoutException.class, () -> second.get(100, TimeUnit.MILLISECONDS));
- release.countDown();
- first.get(5, TimeUnit.SECONDS);
- second.get(5, TimeUnit.SECONDS);
- InOrder order = inOrder(listener);
- order.verify(listener).onParagraphOutputAppend("note", "para", 0, null, "old");
- order.verify(listener).onParagraphOutputUpdated(
- "note", "para", 0, null, InterpreterResult.Type.TEXT, "new");
- } finally {
- release.countDown();
- executor.shutdownNow();
- }
- }
-
- private class BombardEvents implements Runnable {
-
- private final AppendOutputRunner runner;
-
- private BombardEvents(AppendOutputRunner runner) {
- this.runner = runner;
- }
-
- @Override
- public void run() {
- String noteId = "noteId";
- String paraId = "paraId";
- for (int i=0; i log = new ArrayList<>();
-
- @Override
- public boolean requiresLayout() {
- return false;
- }
-
- @Override
- protected void append(final LoggingEvent loggingEvent) {
- log.add(loggingEvent);
- }
-
- @Override
- public void close() {
- }
-
- public List getLog() {
- return new ArrayList<>(log);
- }
- }
-
- private void prepareInvocationCounts(RemoteInterpreterProcessListener listener) {
- doAnswer(new Answer() {
- @Override
- public Void answer(InvocationOnMock invocation) throws Throwable {
- numInvocations += 1;
- return null;
- }
- }).when(listener).onParagraphOutputAppend(
- any(String.class), any(String.class), anyInt(), isNull(), any(String.class));
- }
-
- private void loopForCompletingEvents(RemoteInterpreterProcessListener listener,
- int numTimes, String[][] buffer) {
- numInvocations = 0;
- prepareInvocationCounts(listener);
- AppendOutputRunner runner = new AppendOutputRunner(listener);
- for (String[] bufferElement: buffer) {
- runner.appendBuffer(bufferElement[0], bufferElement[1], 0, null, bufferElement[2]);
- }
- future = service.scheduleWithFixedDelay(runner, 0,
- AppendOutputRunner.BUFFER_TIME_MS, TimeUnit.MILLISECONDS);
- long startTimeMs = System.currentTimeMillis();
- while(numInvocations != numTimes) {
- if (System.currentTimeMillis() - startTimeMs > 2000) {
- fail("Buffered events were not sent for 2 seconds");
- }
- }
+ List batch = new ArrayList<>();
+ batch.add(new AppendOutputBuffer("note", "para", 0, null, "discarded"));
+ runner.run(batch, () -> false);
+ batch.clear();
+ verifyNoInteractions(listener);
+ batch.add(new AppendOutputBuffer("note", "para", 0, null, "new"));
+ runner.run(batch);
+ InOrder order = inOrder(listener);
+ order.verify(listener).onParagraphOutputAppend("note", "para", 0, null, "new");
+ order.verifyNoMoreInteractions();
}
}
diff --git a/zeppelin-server/src/test/java/org/apache/zeppelin/interpreter/remote/ParagraphOutputDispatcherTest.java b/zeppelin-server/src/test/java/org/apache/zeppelin/interpreter/remote/ParagraphOutputDispatcherTest.java
new file mode 100644
index 00000000000..8a58a396ae1
--- /dev/null
+++ b/zeppelin-server/src/test/java/org/apache/zeppelin/interpreter/remote/ParagraphOutputDispatcherTest.java
@@ -0,0 +1,791 @@
+/*
+ * 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.zeppelin.interpreter.remote;
+
+import static org.awaitility.Awaitility.await;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertSame;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTimeoutPreemptively;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.ArgumentMatchers.isNull;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.inOrder;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoInteractions;
+import static org.mockito.Mockito.when;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
+import org.mockito.InOrder;
+import java.lang.management.ManagementFactory;
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
+import org.apache.zeppelin.conf.ZeppelinConfiguration;
+import org.apache.zeppelin.interpreter.InterpreterResult;
+import org.apache.zeppelin.interpreter.InterpreterResultMessage;
+import org.apache.zeppelin.notebook.Note;
+import org.apache.zeppelin.notebook.NoteManager;
+import org.apache.zeppelin.notebook.repo.NotebookRepo;
+import org.apache.zeppelin.user.AuthenticationInfo;
+
+class ParagraphOutputDispatcherTest {
+ @Test
+ void boundariesFlushEarlierAppendsAndKeepLaterAppendsInOrder() throws Exception {
+ RemoteInterpreterProcessListener listener = mock(RemoteInterpreterProcessListener.class);
+ try (ParagraphOutputDispatcher dispatcher = new ParagraphOutputDispatcher(listener)) {
+ dispatcher.appendOutput("note", "para", 0, null, "a");
+ dispatcher.appendOutput("note", "para", 0, null, "b");
+ verifyNoInteractions(listener);
+ dispatcher.updateOutput("note", "para", 0, null, InterpreterResult.Type.TEXT, "replacement")
+ .get(5, TimeUnit.SECONDS);
+ dispatcher.appendOutput("note", "para", 0, null, "after");
+ dispatcher.checkpointOutput("note", "para").get(5, TimeUnit.SECONDS);
+ InOrder order = inOrder(listener);
+ order.verify(listener).onParagraphOutputAppend("note", "para", 0, null, "ab");
+ order.verify(listener).onParagraphOutputUpdated("note", "para", 0, null, InterpreterResult.Type.TEXT, "replacement");
+ order.verify(listener).onParagraphOutputAppend("note", "para", 0, null, "after");
+ order.verify(listener).checkpointOutput("note", "para");
+ order.verifyNoMoreInteractions();
+ }
+ }
+
+ @Test
+ void interleavedAppendsMergeWithinEachBoundary() throws Exception {
+ RemoteInterpreterProcessListener listener = mock(RemoteInterpreterProcessListener.class);
+ try (ParagraphOutputDispatcher dispatcher = new ParagraphOutputDispatcher(listener)) {
+ dispatcher.appendOutput("note", "first", 0, null, "a");
+ dispatcher.appendOutput("note", "second", 0, null, "b");
+ dispatcher.appendOutput("note", "first", 0, null, "c");
+ dispatcher.updateOutput("note", "first", 0, null, InterpreterResult.Type.TEXT, "replace")
+ .get(5, TimeUnit.SECONDS);
+ dispatcher.appendOutput("note", "first", 0, null, "after");
+ dispatcher.checkpointOutput("note", "first").get(5, TimeUnit.SECONDS);
+
+ InOrder order = inOrder(listener);
+ order.verify(listener).onParagraphOutputAppend("note", "first", 0, null, "ac");
+ order.verify(listener).onParagraphOutputAppend("note", "second", 0, null, "b");
+ order.verify(listener).onParagraphOutputUpdated("note", "first", 0, null, InterpreterResult.Type.TEXT, "replace");
+ order.verify(listener).onParagraphOutputAppend("note", "first", 0, null, "after");
+ order.verify(listener).checkpointOutput("note", "first");
+ order.verifyNoMoreInteractions();
+ }
+ }
+
+ @Test
+ void heavilyInterleavedParagraphsUseOneCallbackPerOutputKey() throws Exception {
+ RemoteInterpreterProcessListener listener = mock(RemoteInterpreterProcessListener.class);
+ StringBuilder expected = new StringBuilder();
+ try (ParagraphOutputDispatcher dispatcher = new ParagraphOutputDispatcher(listener, 1, 1000)) {
+ for (int i = 0; i < 250; i++) {
+ String token = i + ";";
+ expected.append(token);
+ for (int paragraph = 0; paragraph < 4; paragraph++) {
+ dispatcher.appendOutput("note", "p" + paragraph, 0, null, token);
+ }
+ }
+ dispatcher.checkpointOutput("note", "p0").get(5, TimeUnit.SECONDS);
+
+ InOrder order = inOrder(listener);
+ for (int paragraph = 0; paragraph < 4; paragraph++) {
+ order.verify(listener).onParagraphOutputAppend("note", "p" + paragraph, 0, null, expected.toString());
+ }
+ order.verify(listener).checkpointOutput("note", "p0");
+ order.verifyNoMoreInteractions();
+ }
+ }
+
+ @Test
+ void scheduledFlushDeliversAppendsWithoutRpcBoundaries() throws Exception {
+ RemoteInterpreterProcessListener listener = mock(RemoteInterpreterProcessListener.class);
+ CountDownLatch delivered = new CountDownLatch(1);
+ doAnswer(call -> {
+ delivered.countDown();
+ return null;
+ }).when(listener).onParagraphOutputAppend("note", "para", 0, null, "periodic");
+ ScheduledExecutorService timer = Executors.newSingleThreadScheduledExecutor();
+ try (ParagraphOutputDispatcher dispatcher = new ParagraphOutputDispatcher(listener)) {
+ dispatcher.appendOutput("note", "para", 0, null, "periodic");
+ timer.scheduleWithFixedDelay(dispatcher::flush, 0,
+ AppendOutputRunner.BUFFER_TIME_MS, TimeUnit.MILLISECONDS);
+ assertTrue(delivered.await(5, TimeUnit.SECONDS));
+ dispatcher.checkpointOutput("note", "para").get(5, TimeUnit.SECONDS);
+ verify(listener).onParagraphOutputAppend("note", "para", 0, null, "periodic");
+ } finally {
+ timer.shutdownNow();
+ }
+ }
+
+ @Test
+ void flushAndSubmissionDoNotWaitForCallbacksOrLoseAnInFlightFlushRequest() throws Exception {
+ RemoteInterpreterProcessListener listener = mock(RemoteInterpreterProcessListener.class);
+ CountDownLatch entered = new CountDownLatch(1);
+ CountDownLatch release = new CountDownLatch(1);
+ CountDownLatch later = new CountDownLatch(1);
+ doAnswer(call -> {
+ entered.countDown();
+ assertTrue(release.await(5, TimeUnit.SECONDS));
+ return null;
+ }).when(listener).onParagraphOutputAppend("note", "para", 0, null, "slow");
+ doAnswer(call -> {
+ later.countDown();
+ return null;
+ }).when(listener).onParagraphOutputAppend("note", "para", 0, null, "later");
+ try (ParagraphOutputDispatcher dispatcher = new ParagraphOutputDispatcher(listener)) {
+ dispatcher.appendOutput("note", "para", 0, null, "slow");
+ dispatcher.flush();
+ assertTrue(entered.await(5, TimeUnit.SECONDS));
+ assertTimeoutPreemptively(Duration.ofSeconds(1), () -> {
+ dispatcher.appendOutput("note", "para", 0, null, "later");
+ dispatcher.flush();
+ });
+ release.countDown();
+ // No further tick or boundary may rescue a lost flush request.
+ assertTrue(later.await(5, TimeUnit.SECONDS));
+ } finally {
+ release.countDown();
+ }
+ }
+
+ @Test
+ void boundaryFailureIsReportedAndDoesNotDiscardLaterEvents() throws Exception {
+ RemoteInterpreterProcessListener listener = mock(RemoteInterpreterProcessListener.class);
+ doThrow(new IllegalStateException("removed")).when(listener)
+ .onParagraphOutputUpdated("note", "gone", 0, null, InterpreterResult.Type.TEXT, "bad");
+ try (ParagraphOutputDispatcher dispatcher = new ParagraphOutputDispatcher(listener)) {
+ Future failed = dispatcher.updateOutput("note", "gone", 0, null, InterpreterResult.Type.TEXT, "bad");
+ ExecutionException failure = assertThrows(ExecutionException.class,
+ () -> failed.get(5, TimeUnit.SECONDS));
+ assertEquals("removed", failure.getCause().getMessage());
+ dispatcher.appendOutput("note", "present", 0, null, "good");
+ dispatcher.checkpointOutput("note", "present").get(5, TimeUnit.SECONDS);
+ verify(listener).onParagraphOutputAppend("note", "present", 0, null, "good");
+ }
+ }
+
+ @ParameterizedTest
+ @ValueSource(strings = {"update", "checkpoint"})
+ void appendErrorFailsBatchBoundariesAndKeepsTheWorker(String boundary) throws Exception {
+ RemoteInterpreterProcessListener listener = mock(RemoteInterpreterProcessListener.class);
+ StackOverflowError error = new StackOverflowError("append failed");
+ doThrow(error).when(listener).onParagraphOutputAppend("note", "para", 0, null, "bad");
+ // With a single worker, the unrelated note progresses only if the worker survives.
+ try (ParagraphOutputDispatcher dispatcher = new ParagraphOutputDispatcher(listener, 1)) {
+ dispatcher.appendOutput("note", "para", 0, null, "bad");
+ Future sameBatch = "update".equals(boundary)
+ ? dispatcher.updateOutput("note", "para", 0, null, InterpreterResult.Type.TEXT, "update")
+ : dispatcher.checkpointOutput("note", "para");
+ ExecutionException failure = assertThrows(ExecutionException.class,
+ () -> sameBatch.get(5, TimeUnit.SECONDS));
+ assertSame(error, failure.getCause());
+ if ("checkpoint".equals(boundary)) {
+ verify(listener, never()).checkpointOutput("note", "para");
+ }
+
+ dispatcher.updateOutput("other", "para", 0, null, InterpreterResult.Type.TEXT, "later")
+ .get(5, TimeUnit.SECONDS);
+ verify(listener).onParagraphOutputUpdated(
+ "other", "para", 0, null, InterpreterResult.Type.TEXT, "later");
+ }
+ }
+
+ @Test
+ void appendErrorReschedulesEventsQueuedBehindTheFailedBatch() throws Exception {
+ RemoteInterpreterProcessListener listener = mock(RemoteInterpreterProcessListener.class);
+ CountDownLatch entered = new CountDownLatch(1);
+ CountDownLatch release = new CountDownLatch(1);
+ doAnswer(call -> {
+ entered.countDown();
+ awaitIgnoringInterrupt(release);
+ return null;
+ }).when(listener).onParagraphOutputAppend("blocker", "para", 0, null, "hold");
+ doThrow(new StackOverflowError("append failed")).when(listener)
+ .onParagraphOutputAppend("note", "para", 0, null, "bad");
+ try (ParagraphOutputDispatcher dispatcher = new ParagraphOutputDispatcher(listener, 1)) {
+ dispatcher.appendOutput("blocker", "para", 0, null, "hold");
+ dispatcher.flush();
+ assertTrue(entered.await(5, TimeUnit.SECONDS));
+
+ // The update queued behind the failed batch must run without waiting for flush().
+ dispatcher.appendOutput("note", "para", 0, null, "bad");
+ Future checkpoint = dispatcher.checkpointOutput("note", "para");
+ Future after =
+ dispatcher.updateOutput("note", "para", 0, null, InterpreterResult.Type.TEXT, "after");
+ release.countDown();
+
+ assertThrows(ExecutionException.class, () -> checkpoint.get(5, TimeUnit.SECONDS));
+ after.get(5, TimeUnit.SECONDS);
+ verify(listener).onParagraphOutputUpdated(
+ "note", "para", 0, null, InterpreterResult.Type.TEXT, "after");
+ } finally {
+ release.countDown();
+ }
+ }
+
+ @Test
+ void largeNoteYieldsAndKeepsAllOutputAcrossWorkerBatches() throws Exception {
+ RemoteInterpreterProcessListener listener = mock(RemoteInterpreterProcessListener.class);
+ CountDownLatch entered = new CountDownLatch(1);
+ CountDownLatch release = new CountDownLatch(1);
+ StringBuilder received = new StringBuilder();
+ AtomicInteger callbacks = new AtomicInteger();
+ AtomicInteger sizeWhenSmallRan = new AtomicInteger();
+ doAnswer(call -> {
+ if ("large".equals(call.getArgument(0))) {
+ received.append((String) call.getArgument(4));
+ if (callbacks.incrementAndGet() == 1) {
+ entered.countDown();
+ awaitIgnoringInterrupt(release);
+ }
+ } else {
+ sizeWhenSmallRan.set(received.length());
+ }
+ return null;
+ }).when(listener).onParagraphOutputAppend(
+ anyString(), anyString(), anyInt(), isNull(), anyString());
+ try (ParagraphOutputDispatcher dispatcher = new ParagraphOutputDispatcher(listener, 1, 2)) {
+ for (int i = 0; i < 5; i++) {
+ dispatcher.appendOutput("large", "para", 0, null, "x");
+ }
+ Future largeDone = dispatcher.checkpointOutput("large", "para");
+ assertTrue(entered.await(5, TimeUnit.SECONDS));
+ // Submit only after a large batch has been claimed, rather than ahead of its ready entry.
+ dispatcher.appendOutput("small", "para", 0, null, "small");
+ Future smallDone = dispatcher.checkpointOutput("small", "para");
+ release.countDown();
+ smallDone.get(5, TimeUnit.SECONDS);
+ largeDone.get(5, TimeUnit.SECONDS);
+ assertTrue(sizeWhenSmallRan.get() > 0);
+ assertTrue(sizeWhenSmallRan.get() < 5, "Small note must run before all large output");
+ InOrder order = inOrder(listener);
+ order.verify(listener).onParagraphOutputAppend("small", "para", 0, null, "small");
+ order.verify(listener).checkpointOutput("large", "para");
+ assertEquals("x".repeat(5), received.toString());
+ assertTrue(callbacks.get() < 5, "Adjacent appends must still be batched");
+ } finally {
+ release.countDown();
+ }
+ }
+
+ @Test
+ void concurrentProducersKeepOneWriterAndDeliverEveryEventOnceInProducerOrder() throws Exception {
+ RemoteInterpreterProcessListener listener = mock(RemoteInterpreterProcessListener.class);
+ AtomicInteger active = new AtomicInteger();
+ AtomicInteger maximum = new AtomicInteger();
+ List received = Collections.synchronizedList(new ArrayList<>());
+ doAnswer(call -> {
+ maximum.accumulateAndGet(active.incrementAndGet(), Math::max);
+ try {
+ String output = call.getArgument(4);
+ Collections.addAll(received, output.split("\n"));
+ } finally {
+ active.decrementAndGet();
+ }
+ return null;
+ }).when(listener).onParagraphOutputAppend(
+ anyString(), anyString(), anyInt(), isNull(), anyString());
+ ExecutorService producers = Executors.newFixedThreadPool(4);
+ ScheduledExecutorService timer = Executors.newSingleThreadScheduledExecutor();
+ try (ParagraphOutputDispatcher dispatcher = new ParagraphOutputDispatcher(listener)) {
+ timer.scheduleWithFixedDelay(dispatcher::flush, 0, 1, TimeUnit.MILLISECONDS);
+ List> futures = new ArrayList<>();
+ for (int producer = 0; producer < 4; producer++) {
+ final int id = producer;
+ futures.add(producers.submit(() -> {
+ for (int i = 0; i < 1000; i++) {
+ dispatcher.appendOutput("note", "para", 0, null, id + ":" + i + "\n");
+ if (i % 100 == 0) {
+ dispatcher.checkpointOutput("note", "para").get(5, TimeUnit.SECONDS);
+ }
+ }
+ return null;
+ }));
+ }
+ for (Future future : futures) {
+ future.get(10, TimeUnit.SECONDS);
+ }
+ dispatcher.checkpointOutput("note", "para").get(5, TimeUnit.SECONDS);
+ assertEquals(1, maximum.get());
+ assertEquals(4000, received.size());
+ assertEquals(4000, new HashSet<>(received).size());
+ int[] next = new int[4];
+ for (String token : received) {
+ String[] parts = token.split(":");
+ int producer = Integer.parseInt(parts[0]);
+ assertEquals(next[producer]++, Integer.parseInt(parts[1]));
+ }
+ } finally {
+ timer.shutdownNow();
+ producers.shutdownNow();
+ }
+ }
+
+ @Test
+ void blockedCheckpointsDoNotOccupyOutputWorkersOrReleaseTheirNotes() throws Exception {
+ RemoteInterpreterProcessListener listener = mock(RemoteInterpreterProcessListener.class);
+ CountDownLatch entered = new CountDownLatch(2);
+ CountDownLatch release = new CountDownLatch(1);
+ CountDownLatch unrelatedOutput = new CountDownLatch(1);
+ CountDownLatch laterOutput = new CountDownLatch(1);
+ doAnswer(call -> {
+ entered.countDown();
+ awaitIgnoringInterrupt(release);
+ return null;
+ }).when(listener).checkpointOutput(anyString(), anyString());
+ doAnswer(call -> {
+ unrelatedOutput.countDown();
+ return null;
+ }).when(listener).onParagraphOutputAppend("C", "para", 0, null, "unrelated");
+ doAnswer(call -> {
+ laterOutput.countDown();
+ return null;
+ }).when(listener).onParagraphOutputAppend("A", "para", 0, null, "later");
+ try (ParagraphOutputDispatcher dispatcher = new ParagraphOutputDispatcher(listener, 2)) {
+ Future first = dispatcher.checkpointOutput("A", "para");
+ Future second = dispatcher.checkpointOutput("B", "para");
+ assertTrue(entered.await(5, TimeUnit.SECONDS));
+ dispatcher.appendOutput("A", "para", 0, null, "later");
+ dispatcher.appendOutput("C", "para", 0, null, "unrelated");
+ Future third = dispatcher.checkpointOutput("C", "para");
+ dispatcher.flush();
+ assertTrue(unrelatedOutput.await(5, TimeUnit.SECONDS));
+ assertFalse(laterOutput.await(100, TimeUnit.MILLISECONDS));
+ assertFalse(first.isDone());
+ assertFalse(second.isDone());
+ assertFalse(third.isDone());
+ verify(listener, never()).checkpointOutput("C", "para");
+ release.countDown();
+ first.get(5, TimeUnit.SECONDS);
+ second.get(5, TimeUnit.SECONDS);
+ third.get(5, TimeUnit.SECONDS);
+ assertTrue(laterOutput.await(5, TimeUnit.SECONDS));
+ verify(listener).checkpointOutput("A", "para");
+ verify(listener).checkpointOutput("B", "para");
+ InOrder order = inOrder(listener);
+ order.verify(listener).checkpointOutput("A", "para");
+ order.verify(listener).onParagraphOutputAppend("A", "para", 0, null, "later");
+ } finally {
+ release.countDown();
+ }
+ }
+
+ @Test
+ void noteManagerSaveLockDoesNotOccupyOutputWorkers() throws Exception {
+ NotebookRepo repo = mock(NotebookRepo.class);
+ when(repo.list(any())).thenReturn(Collections.emptyMap());
+ NoteManager manager = new NoteManager(repo, ZeppelinConfiguration.load());
+ Note firstNote = new Note();
+ firstNote.setId("A");
+ firstNote.setPath("/A");
+ Note secondNote = new Note();
+ secondNote.setId("B");
+ secondNote.setPath("/B");
+ Note unrelatedNote = new Note();
+ unrelatedNote.setId("C");
+ unrelatedNote.setPath("/C");
+ manager.saveNote(unrelatedNote);
+ CountDownLatch saveStarted = new CountDownLatch(1);
+ CountDownLatch releaseSave = new CountDownLatch(1);
+ AtomicReference savingThread = new AtomicReference<>();
+ doAnswer(call -> {
+ if (call.getArgument(0) == firstNote) {
+ savingThread.set(Thread.currentThread());
+ saveStarted.countDown();
+ awaitIgnoringInterrupt(releaseSave);
+ }
+ return null;
+ }).when(repo).save(any(), any());
+
+ RemoteInterpreterProcessListener listener = mock(RemoteInterpreterProcessListener.class);
+ doAnswer(call -> {
+ manager.saveNote(firstNote);
+ return null;
+ }).when(listener).checkpointOutput("A", "para");
+ CountDownLatch secondStarted = new CountDownLatch(1);
+ AtomicReference secondThread = new AtomicReference<>();
+ doAnswer(call -> {
+ secondThread.set(Thread.currentThread());
+ secondStarted.countDown();
+ manager.saveNote(secondNote);
+ return null;
+ }).when(listener).checkpointOutput("B", "para");
+ CountDownLatch unrelatedOutput = new CountDownLatch(1);
+ doAnswer(call -> {
+ manager.processNote("C", note -> {
+ unrelatedOutput.countDown();
+ return null;
+ });
+ return null;
+ }).when(listener).onParagraphOutputAppend("C", "para", 0, null, "output");
+
+ try (ParagraphOutputDispatcher dispatcher = new ParagraphOutputDispatcher(listener, 2)) {
+ Future first = dispatcher.checkpointOutput("A", "para");
+ assertTrue(saveStarted.await(5, TimeUnit.SECONDS));
+ Future second = dispatcher.checkpointOutput("B", "para");
+ assertTrue(secondStarted.await(5, TimeUnit.SECONDS));
+ await().atMost(Duration.ofSeconds(5)).until(() ->
+ secondThread.get().getState() == Thread.State.BLOCKED
+ && ManagementFactory.getThreadMXBean().getThreadInfo(secondThread.get().getId())
+ .getLockOwnerId() == savingThread.get().getId());
+ dispatcher.appendOutput("C", "para", 0, null, "output");
+ dispatcher.flush();
+ assertTrue(unrelatedOutput.await(5, TimeUnit.SECONDS));
+ assertFalse(first.isDone());
+ assertFalse(second.isDone());
+ releaseSave.countDown();
+ first.get(5, TimeUnit.SECONDS);
+ second.get(5, TimeUnit.SECONDS);
+ verify(repo).save(firstNote, AuthenticationInfo.ANONYMOUS);
+ verify(repo).save(secondNote, AuthenticationInfo.ANONYMOUS);
+ } finally {
+ releaseSave.countDown();
+ }
+ }
+
+ @Test
+ void failedCheckpointReleasesItsNoteForLaterOutput() throws Exception {
+ RemoteInterpreterProcessListener listener = mock(RemoteInterpreterProcessListener.class);
+ doThrow(new IllegalStateException("save failed")).when(listener)
+ .checkpointOutput("note", "para");
+ CountDownLatch delivered = new CountDownLatch(1);
+ doAnswer(call -> {
+ delivered.countDown();
+ return null;
+ }).when(listener).onParagraphOutputAppend("note", "para", 0, null, "after");
+ try (ParagraphOutputDispatcher dispatcher = new ParagraphOutputDispatcher(listener, 1)) {
+ Future checkpoint = dispatcher.checkpointOutput("note", "para");
+ dispatcher.appendOutput("note", "para", 0, null, "after");
+ ExecutionException failure = assertThrows(ExecutionException.class,
+ () -> checkpoint.get(5, TimeUnit.SECONDS));
+ assertEquals("save failed", failure.getCause().getMessage());
+ assertTrue(delivered.await(5, TimeUnit.SECONDS));
+ InOrder order = inOrder(listener);
+ order.verify(listener).checkpointOutput("note", "para");
+ order.verify(listener).onParagraphOutputAppend("note", "para", 0, null, "after");
+ }
+ }
+
+ @Test
+ void checkpointStopsBatchBeforeLaterAppend() throws Exception {
+ RemoteInterpreterProcessListener listener = mock(RemoteInterpreterProcessListener.class);
+ CountDownLatch entered = new CountDownLatch(1);
+ CountDownLatch release = new CountDownLatch(1);
+ CountDownLatch later = new CountDownLatch(1);
+ doAnswer(call -> {
+ entered.countDown();
+ awaitIgnoringInterrupt(release);
+ return null;
+ }).when(listener).onParagraphOutputAppend("blocker", "para", 0, null, "busy");
+ doAnswer(call -> {
+ later.countDown();
+ return null;
+ }).when(listener).onParagraphOutputAppend("note", "para", 0, null, "after");
+ try (ParagraphOutputDispatcher dispatcher = new ParagraphOutputDispatcher(listener, 1)) {
+ dispatcher.appendOutput("blocker", "para", 0, null, "busy");
+ dispatcher.flush();
+ assertTrue(entered.await(5, TimeUnit.SECONDS));
+ dispatcher.appendOutput("note", "para", 0, null, "before");
+ Future checkpoint = dispatcher.checkpointOutput("note", "para");
+ dispatcher.appendOutput("note", "para", 0, null, "after");
+ dispatcher.flush();
+ release.countDown();
+ checkpoint.get(5, TimeUnit.SECONDS);
+ assertTrue(later.await(5, TimeUnit.SECONDS));
+ InOrder order = inOrder(listener);
+ order.verify(listener).onParagraphOutputAppend("note", "para", 0, null, "before");
+ order.verify(listener).checkpointOutput("note", "para");
+ order.verify(listener).onParagraphOutputAppend("note", "para", 0, null, "after");
+ } finally {
+ release.countDown();
+ }
+ }
+
+ @Test
+ void idleQueuesAreReclaimedAndTheSameNoteCanDeliverAgain() throws Exception {
+ RemoteInterpreterProcessListener listener = mock(RemoteInterpreterProcessListener.class);
+ try (ParagraphOutputDispatcher dispatcher = new ParagraphOutputDispatcher(listener)) {
+ for (int i = 0; i < 200; i++) {
+ dispatcher.appendOutput("note" + i, "para", 0, null, "data");
+ dispatcher.checkpointOutput("note" + i, "para").get(5, TimeUnit.SECONDS);
+ }
+ await().atMost(Duration.ofSeconds(5)).until(() -> dispatcher.pendingNoteCount() == 0);
+ dispatcher.appendOutput("note0", "para", 0, null, "again");
+ dispatcher.checkpointOutput("note0", "para").get(5, TimeUnit.SECONDS);
+ verify(listener).onParagraphOutputAppend("note0", "para", 0, null, "data");
+ verify(listener).onParagraphOutputAppend("note0", "para", 0, null, "again");
+ await().atMost(Duration.ofSeconds(5)).until(() -> dispatcher.pendingNoteCount() == 0);
+ }
+ }
+
+ @Test
+ void closeCancelsABoundaryAlreadyDrainedBehindABlockedAppend() throws Exception {
+ RemoteInterpreterProcessListener listener = mock(RemoteInterpreterProcessListener.class);
+ CountDownLatch entered = new CountDownLatch(1);
+ CountDownLatch release = new CountDownLatch(1);
+ AtomicReference callbackThread = new AtomicReference<>();
+ doAnswer(call -> {
+ callbackThread.set(Thread.currentThread());
+ entered.countDown();
+ try {
+ release.await(5, TimeUnit.SECONDS);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ }
+ return null;
+ }).when(listener).onParagraphOutputAppend("note", "para", 0, null, "blocked");
+ try (ParagraphOutputDispatcher dispatcher = new ParagraphOutputDispatcher(listener, 1)) {
+ dispatcher.appendOutput("note", "para", 0, null, "blocked");
+ Future boundary = dispatcher.checkpointOutput("note", "para");
+ assertTrue(entered.await(5, TimeUnit.SECONDS));
+ dispatcher.close();
+ assertStopped(boundary);
+ release.countDown();
+ callbackThread.get().join(5000);
+ assertFalse(callbackThread.get().isAlive());
+ verify(listener, never()).checkpointOutput("note", "para");
+ } finally {
+ release.countDown();
+ }
+ }
+
+ @Test
+ void closeFailsQueuedAndInFlightBoundariesAndRejectsNewEvents() throws Exception {
+ RemoteInterpreterProcessListener listener = mock(RemoteInterpreterProcessListener.class);
+ CountDownLatch entered = new CountDownLatch(1);
+ CountDownLatch release = new CountDownLatch(1);
+ doAnswer(call -> {
+ entered.countDown();
+ try {
+ release.await(5, TimeUnit.SECONDS);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ }
+ return null;
+ }).when(listener).checkpointOutput("A", "para");
+ try (ParagraphOutputDispatcher dispatcher = new ParagraphOutputDispatcher(listener, 1)) {
+ Future inFlight = dispatcher.checkpointOutput("A", "para");
+ assertTrue(entered.await(5, TimeUnit.SECONDS));
+ Future queued = dispatcher.checkpointOutput("B", "para");
+ dispatcher.close();
+ assertStopped(inFlight);
+ assertStopped(queued);
+ assertEquals(0, dispatcher.pendingNoteCount());
+ assertThrows(IllegalStateException.class, () ->
+ dispatcher.appendOutput("A", "para", 0, null, "late"));
+ assertThrows(IllegalStateException.class, () -> dispatcher.checkpointOutput("A", "para"));
+ verify(listener, never()).checkpointOutput("B", "para");
+ } finally {
+ release.countDown();
+ }
+ }
+
+ @ParameterizedTest
+ @ValueSource(booleans = {false, true})
+ void shutdownStopsLaterAppendGroupsEvenWhenActiveCallbackIgnoresInterrupt(boolean withBoundary)
+ throws Exception {
+ RemoteInterpreterProcessListener listener = mock(RemoteInterpreterProcessListener.class);
+ CountDownLatch entered = new CountDownLatch(1);
+ CountDownLatch release = new CountDownLatch(1);
+ AtomicReference worker = new AtomicReference<>();
+ doAnswer(call -> {
+ worker.set(Thread.currentThread());
+ entered.countDown();
+ awaitIgnoringInterrupt(release);
+ return null;
+ }).when(listener).onParagraphOutputAppend("note", "first", 0, null, "a");
+ try (ParagraphOutputDispatcher dispatcher = new ParagraphOutputDispatcher(listener, 1)) {
+ dispatcher.appendOutput("note", "first", 0, null, "a");
+ dispatcher.appendOutput("note", "second", 0, null, "b");
+ Future boundary = null;
+ if (withBoundary) {
+ boundary = dispatcher.checkpointOutput("note", "para");
+ } else {
+ dispatcher.flush();
+ }
+ assertTrue(entered.await(5, TimeUnit.SECONDS));
+ dispatcher.close();
+ if (boundary != null) {
+ assertStopped(boundary);
+ }
+ release.countDown();
+ worker.get().join(5000);
+ assertFalse(worker.get().isAlive());
+ verify(listener).onParagraphOutputAppend("note", "first", 0, null, "a");
+ verify(listener, never()).onParagraphOutputAppend("note", "second", 0, null, "b");
+ verify(listener, never()).checkpointOutput("note", "para");
+ } finally {
+ release.countDown();
+ }
+ }
+
+ @Test
+ void closeReleasesAllBoundaryWaitersWithoutWaitingForAnActiveCallback() throws Exception {
+ RemoteInterpreterProcessListener listener = mock(RemoteInterpreterProcessListener.class);
+ CountDownLatch entered = new CountDownLatch(1);
+ CountDownLatch release = new CountDownLatch(1);
+ CountDownLatch interrupted = new CountDownLatch(1);
+ AtomicReference worker = new AtomicReference<>();
+ doAnswer(call -> {
+ worker.set(Thread.currentThread());
+ entered.countDown();
+ try {
+ release.await();
+ } catch (InterruptedException e) {
+ interrupted.countDown();
+ awaitIgnoringInterrupt(release);
+ }
+ return null;
+ }).when(listener).checkpointOutput("note", "running");
+ ExecutorService callers = Executors.newFixedThreadPool(4);
+ try (ParagraphOutputDispatcher dispatcher = new ParagraphOutputDispatcher(listener, 1)) {
+ List> boundaries = new ArrayList<>();
+ boundaries.add(dispatcher.checkpointOutput("note", "running"));
+ assertTrue(entered.await(5, TimeUnit.SECONDS));
+ for (int i = 0; i < 20; i++) {
+ boundaries.add(dispatcher.checkpointOutput("note", "queued" + i));
+ boundaries.add(dispatcher.checkpointOutput("other", "queued" + i));
+ }
+ CountDownLatch waiting = new CountDownLatch(4);
+ List> waiters = new ArrayList<>();
+ for (int i = 0; i < 4; i++) {
+ Future boundary = boundaries.get(i);
+ waiters.add(callers.submit(() -> {
+ waiting.countDown();
+ assertStopped(boundary);
+ }));
+ }
+ assertTrue(waiting.await(5, TimeUnit.SECONDS));
+ assertTimeoutPreemptively(Duration.ofSeconds(1), dispatcher::close);
+ assertTrue(interrupted.await(5, TimeUnit.SECONDS));
+ for (Future boundary : boundaries) {
+ assertStopped(boundary);
+ }
+ for (Future waiter : waiters) {
+ waiter.get(5, TimeUnit.SECONDS);
+ }
+ assertEquals(1, release.getCount(), "Callback must still be blocked after waiters exit");
+ assertThrows(IllegalStateException.class,
+ () -> dispatcher.appendOutput("other", "para", 0, null, "late"));
+ verify(listener, never()).checkpointOutput("note", "queued0");
+ verify(listener, never()).checkpointOutput("other", "queued0");
+ release.countDown();
+ worker.get().join(5000);
+ assertFalse(worker.get().isAlive(), "Shutdown must survive a callback clearing interruption");
+ } finally {
+ release.countDown();
+ callers.shutdownNow();
+ }
+ }
+
+ @Test
+ void acceptedBoundaryCannotBeCancelledOrSkipped() throws Exception {
+ RemoteInterpreterProcessListener listener = mock(RemoteInterpreterProcessListener.class);
+ CountDownLatch entered = new CountDownLatch(1);
+ CountDownLatch release = new CountDownLatch(1);
+ doAnswer(call -> {
+ entered.countDown();
+ awaitIgnoringInterrupt(release);
+ return null;
+ }).when(listener).onParagraphOutputAppend("note", "para", 0, null, "before");
+ try (ParagraphOutputDispatcher dispatcher = new ParagraphOutputDispatcher(listener, 1)) {
+ dispatcher.appendOutput("note", "para", 0, null, "before");
+ Future boundary = dispatcher.checkpointOutput("note", "para");
+ assertTrue(entered.await(5, TimeUnit.SECONDS));
+ assertFalse(boundary.cancel(true));
+ assertFalse(boundary.isCancelled());
+ release.countDown();
+ boundary.get(5, TimeUnit.SECONDS);
+ verify(listener).checkpointOutput("note", "para");
+ } finally {
+ release.countDown();
+ }
+ }
+
+ @Test
+ void queuedUpdateAllUsesASnapshotOfItsReplacementList() throws Exception {
+ RemoteInterpreterProcessListener listener = mock(RemoteInterpreterProcessListener.class);
+ CountDownLatch entered = new CountDownLatch(1);
+ CountDownLatch release = new CountDownLatch(1);
+ doAnswer(call -> {
+ entered.countDown();
+ assertTrue(release.await(5, TimeUnit.SECONDS));
+ return null;
+ }).when(listener).onParagraphOutputAppend("note", "para", 0, null, "old");
+ try (ParagraphOutputDispatcher dispatcher = new ParagraphOutputDispatcher(listener, 1)) {
+ dispatcher.appendOutput("note", "para", 0, null, "old");
+ dispatcher.flush();
+ assertTrue(entered.await(5, TimeUnit.SECONDS));
+ List replacements = new ArrayList<>();
+ replacements.add(new InterpreterResultMessage(InterpreterResult.Type.TEXT, "new"));
+ Future updated = dispatcher.updateAllOutput("note", "para", null, replacements);
+ replacements.clear();
+ release.countDown();
+ updated.get(5, TimeUnit.SECONDS);
+ InOrder order = inOrder(listener);
+ order.verify(listener).onParagraphOutputAppend("note", "para", 0, null, "old");
+ order.verify(listener).onParagraphOutputClear("note", "para", null);
+ order.verify(listener).onParagraphOutputUpdated("note", "para", 0, null, InterpreterResult.Type.TEXT, "new");
+ order.verifyNoMoreInteractions();
+ } finally {
+ release.countDown();
+ }
+ }
+
+ private void assertStopped(Future completion) {
+ ExecutionException failure = assertThrows(ExecutionException.class,
+ () -> completion.get(5, TimeUnit.SECONDS));
+ assertTrue(failure.getCause() instanceof IllegalStateException);
+ }
+
+ private void awaitIgnoringInterrupt(CountDownLatch release) {
+ boolean interrupted = false;
+ try {
+ while (true) {
+ try {
+ release.await();
+ return;
+ } catch (InterruptedException e) {
+ interrupted = true;
+ }
+ }
+ } finally {
+ if (interrupted) {
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+}