Skip to content

[ZEPPELIN-6700] Convert a broadcast message to JSON once instead of once per connection #5462

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
big-cir wants to merge 1 commit into
base: master
Choose a base branch
from
+120 −10
Open

[ZEPPELIN-6700] Convert a broadcast message to JSON once instead of once per connection #5462

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 22 additions & 10 deletions zeppelin-server/src/main/java/org/apache/zeppelin/socket/ConnectionManager.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -229,10 +229,11 @@ protected String serializeMessage(Message m) {
}

public void broadcast(Message m) {
String serializedMsg = serializeMessage(m);
synchronized (connectedSockets) {
for (NotebookSocket ns : connectedSockets) {
try {
ns.send(serializeMessage(m));
ns.send(serializedMsg);
} catch (IOException | RuntimeException e) {
LOGGER.error("Send error: {}", m, e);
}
Expand All @@ -251,9 +252,10 @@ public void broadcast(String noteId, Message m) {
socketsToBroadcast = new ArrayList<>(sockets);
}
LOGGER.debug("SEND >> {}", m);
String serializedMsg = serializeMessage(m);
for (NotebookSocket conn : socketsToBroadcast) {
try {
conn.send(serializeMessage(m));
conn.send(serializedMsg);
} catch (IOException | RuntimeException e) {
LOGGER.error("socket error", e);
}
Expand All @@ -262,14 +264,17 @@ public void broadcast(String noteId, Message m) {

private void broadcastToWatchers(String noteId, String subject, Message message) {
synchronized (watcherSockets) {
if (watcherSockets.isEmpty()) {
return;
}
String watcherMsg = WatcherMessage.builder(noteId)
.subject(subject)
.message(serializeMessage(message))
.build()
.toJson();
for (NotebookSocket watcher : watcherSockets) {
try {
watcher.send(
WatcherMessage.builder(noteId)
.subject(subject)
.message(serializeMessage(message))
.build()
.toJson());
watcher.send(watcherMsg);
} catch (IOException | RuntimeException e) {
LOGGER.error("Cannot broadcast message to watcher", e);
}
Expand All @@ -289,12 +294,13 @@ public void broadcastExcept(String noteId, Message m, NotebookSocket exclude) {
}

LOGGER.debug("SEND >> {}", m);
String serializedMsg = serializeMessage(m);
for (NotebookSocket conn : socketsToBroadcast) {
if (exclude.equals(conn)) {
continue;
}
try {
conn.send(serializeMessage(m));
conn.send(serializedMsg);
} catch (IOException | RuntimeException e) {
LOGGER.error("socket error", e);
}
Expand Down Expand Up @@ -340,8 +346,14 @@ public void multicastToUser(String user, Message m) {
return;
}

String serializedMsg = serializeMessage(m);
for (NotebookSocket conn : connections) {
unicast(m, conn);
try {
conn.send(serializedMsg);
} catch (IOException | RuntimeException e) {
LOGGER.error("socket error", e);
}
broadcastToWatchers(StringUtils.EMPTY, StringUtils.EMPTY, m);
}
}

Expand Down
98 changes: 98 additions & 0 deletions zeppelin-server/src/test/java/org/apache/zeppelin/socket/ConnectionManagerTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,9 +22,13 @@
import static org.junit.jupiter.api.Assertions.assertNull;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.Mockito.clearInvocations;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

import java.io.IOException;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
Expand All @@ -34,12 +38,15 @@
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;

import org.apache.zeppelin.common.Message;
import org.apache.zeppelin.conf.ZeppelinConfiguration;
import org.apache.zeppelin.notebook.AuthorizationService;
import org.apache.zeppelin.util.WatcherSecurityKey;
import org.junit.jupiter.api.Test;
import org.mockito.ArgumentCaptor;

class ConnectionManagerTest {

Expand Down Expand Up @@ -310,4 +317,95 @@ void removeUserConnectionBeforeUserAssignment() {
assertDoesNotThrow(() -> manager.removeUserConnection("", socket));
assertTrue(manager.userSocketMap.isEmpty());
}

@Test
void broadcastToNoteSerializesMessageOnce() throws IOException {
CountingConnectionManager manager = newCountingManager();
NotebookSocket first = mock(NotebookSocket.class);
NotebookSocket second = mock(NotebookSocket.class);
NotebookSocket third = mock(NotebookSocket.class);
manager.addNoteConnection("note1", first);
manager.addNoteConnection("note1", second);
manager.addNoteConnection("note1", third);
// adding a connection broadcasts the collaborative mode status, which is not under test
manager.serializeCount.set(0);
clearInvocations(first, second, third);

manager.broadcast("note1", new Message(Message.OP.NOTE).put("note", "payload"));

assertEquals(1, manager.serializeCount.get());
assertEquals(payloadSentTo(first), payloadSentTo(second));
assertEquals(payloadSentTo(first), payloadSentTo(third));
}

@Test
void broadcastExceptSerializesMessageOnce() throws IOException {
CountingConnectionManager manager = newCountingManager();
NotebookSocket sender = mock(NotebookSocket.class);
NotebookSocket first = mock(NotebookSocket.class);
NotebookSocket second = mock(NotebookSocket.class);
manager.addNoteConnection("note1", sender);
manager.addNoteConnection("note1", first);
manager.addNoteConnection("note1", second);
manager.serializeCount.set(0);
clearInvocations(sender, first, second);

manager.broadcastExcept("note1",
new Message(Message.OP.PATCH_PARAGRAPH).put("patch", "payload"), sender);

assertEquals(1, manager.serializeCount.get());
assertEquals(payloadSentTo(first), payloadSentTo(second));
verify(sender, times(0)).send(anyString());
}

@Test
void multicastToUserSerializesMessageOnceAndKeepsWatcherBroadcasts() throws IOException {
CountingConnectionManager manager = newCountingManager();
NotebookSocket first = mock(NotebookSocket.class);
NotebookSocket second = mock(NotebookSocket.class);
manager.addUserConnection("testUser", first);
manager.addUserConnection("testUser", second);
NotebookSocket watcher = mock(NotebookSocket.class);
when(watcher.getHeader(WatcherSecurityKey.HTTP_HEADER)).thenReturn(WatcherSecurityKey.getKey());
manager.switchConnectionToWatcher(watcher);
manager.serializeCount.set(0);
clearInvocations(first, second, watcher);

manager.multicastToUser("testUser", new Message(Message.OP.NOTES_INFO).put("notes", "payload"));

assertEquals(payloadSentTo(first), payloadSentTo(second));
// one conversion for the two connections, plus one per watcher broadcast as before
assertEquals(3, manager.serializeCount.get());
verify(watcher, times(2)).send(anyString());
}

private static CountingConnectionManager newCountingManager() {
return new CountingConnectionManager(mock(AuthorizationService.class),
ZeppelinConfiguration.load());
}

private static String payloadSentTo(NotebookSocket socket) throws IOException {
ArgumentCaptor<String> payload = ArgumentCaptor.forClass(String.class);
verify(socket).send(payload.capture());
return payload.getValue();
}

/**
* Counts how often a broadcast converts its message, which is what distinguishes a single
* conversion reused across connections from one conversion per connection.
*/
private static class CountingConnectionManager extends ConnectionManager {
private final AtomicInteger serializeCount = new AtomicInteger();

CountingConnectionManager(AuthorizationService authorizationService,
ZeppelinConfiguration zConf) {
super(authorizationService, zConf);
}

@Override
protected String serializeMessage(Message m) {
serializeCount.incrementAndGet();
return super.serializeMessage(m);
}
}
}
Loading