mirror of
https://github.com/apache/zeppelin
synced 2026-05-24 09:38:26 +00:00
rework watcher creation and ws session with notes
This commit is contained in:
Anthony Corbacho
parent
3bd34827fb
commit
c32192ade2
1 changed files with 61 additions and 80 deletions
|
|
@ -20,6 +20,8 @@ import java.io.IOException;
|
|||
import java.net.URI;
|
||||
import java.util.HashMap;
|
||||
import java.util.Map;
|
||||
import java.util.Timer;
|
||||
import java.util.concurrent.ConcurrentHashMap;
|
||||
import java.util.concurrent.ExecutionException;
|
||||
import java.util.concurrent.Future;
|
||||
|
||||
|
|
@ -27,6 +29,7 @@ import org.apache.commons.lang3.StringUtils;
|
|||
import org.apache.zeppelin.conf.ZeppelinConfiguration;
|
||||
import org.apache.zeppelin.notebook.repo.zeppelinhub.security.Authentication;
|
||||
import org.apache.zeppelin.notebook.repo.zeppelinhub.websocket.listener.WatcherWebsocket;
|
||||
import org.apache.zeppelin.notebook.repo.zeppelinhub.websocket.listener.ZeppelinWebsocket;
|
||||
import org.apache.zeppelin.notebook.repo.zeppelinhub.websocket.protocol.ZeppelinhubMessage;
|
||||
import org.apache.zeppelin.notebook.repo.zeppelinhub.websocket.scheduler.SchedulerService;
|
||||
import org.apache.zeppelin.notebook.repo.zeppelinhub.websocket.scheduler.ZeppelinHeartbeat;
|
||||
|
|
@ -52,12 +55,14 @@ public class ZeppelinClient {
|
|||
private final String zeppelinhubToken;
|
||||
private final WebSocketClient wsClient;
|
||||
private static Gson gson;
|
||||
//private ConcurrentHashMap<String, Session> zeppelinConnectionMap;
|
||||
// Keep track of current open connection per notebook.
|
||||
private ConcurrentHashMap<String, Session> notesConnection;
|
||||
// Listen to every note actions.
|
||||
private static Session watcherSession;
|
||||
private static ZeppelinClient instance = null;
|
||||
private SchedulerService schedulerService;
|
||||
private Authentication authModule;
|
||||
private static final int min = 60;
|
||||
private static final int MIN = 60;
|
||||
|
||||
public static ZeppelinClient initialize(String zeppelinUrl, String token,
|
||||
ZeppelinConfiguration conf) {
|
||||
|
|
@ -76,20 +81,19 @@ public class ZeppelinClient {
|
|||
zeppelinhubToken = token;
|
||||
wsClient = createNewWebsocketClient();
|
||||
gson = new Gson();
|
||||
//zeppelinConnectionMap = new ConcurrentHashMap<>();
|
||||
notesConnection = new ConcurrentHashMap<>();
|
||||
schedulerService = SchedulerService.getInstance();
|
||||
authModule = Authentication.initialize(token, conf);
|
||||
if (authModule != null) {
|
||||
SchedulerService.getInstance().addOnce(authModule, 10);
|
||||
}
|
||||
//
|
||||
LOG.info("Initialized Zeppelin websocket client on {}", zeppelinWebsocketUrl);
|
||||
}
|
||||
|
||||
private WebSocketClient createNewWebsocketClient() {
|
||||
SslContextFactory sslContextFactory = new SslContextFactory();
|
||||
WebSocketClient client = new WebSocketClient(sslContextFactory);
|
||||
client.setMaxIdleTimeout(5 * min * 1000);
|
||||
client.setMaxIdleTimeout(5 * MIN * 1000);
|
||||
client.setMaxTextMessageBufferSize(Client.getMaxNoteSize());
|
||||
client.getPolicy().setMaxTextMessageSize(Client.getMaxNoteSize());
|
||||
//TODO(khalid): other client settings
|
||||
|
|
@ -110,11 +114,11 @@ public class ZeppelinClient {
|
|||
}
|
||||
|
||||
private void addRoutines() {
|
||||
schedulerService.add(ZeppelinHeartbeat.newInstance(this), 15, 4 * min);
|
||||
new java.util.Timer().schedule(new java.util.TimerTask() {
|
||||
schedulerService.add(ZeppelinHeartbeat.newInstance(this), 10, 1 * MIN);
|
||||
new Timer().schedule(new java.util.TimerTask() {
|
||||
@Override
|
||||
public void run() {
|
||||
openWatcherSession();
|
||||
watcherSession = openWatcherSession();
|
||||
}
|
||||
}, 5000);
|
||||
}
|
||||
|
|
@ -122,11 +126,13 @@ public class ZeppelinClient {
|
|||
public void stop() {
|
||||
try {
|
||||
if (wsClient != null) {
|
||||
//removeAllZeppelinConnections();
|
||||
wsClient.stop();
|
||||
} else {
|
||||
LOG.warn("Cannot stop zeppelin websocket client - isn't initialized");
|
||||
}
|
||||
if (watcherSession != null) {
|
||||
watcherSession.close();
|
||||
}
|
||||
} catch (Exception e) {
|
||||
LOG.error("Cannot stop Zeppelin websocket client", e);
|
||||
}
|
||||
|
|
@ -176,34 +182,23 @@ public class ZeppelinClient {
|
|||
}
|
||||
|
||||
public void send(Message msg, String noteId) {
|
||||
//Session noteSession = getZeppelinConnection(noteId);
|
||||
//if (!isSessionOpen(noteSession)) {
|
||||
// LOG.error("Cannot open websocket connection to Zeppelin note {}", noteId);
|
||||
// return;
|
||||
// }
|
||||
// noteSession.getRemote().sendStringByFuture(serialize(msg));
|
||||
}
|
||||
|
||||
private boolean isSessionOpen(Session session) {
|
||||
return (session != null) && (session.isOpen());
|
||||
Session noteSession = getZeppelinConnection(noteId);
|
||||
if (!isSessionOpen(noteSession)) {
|
||||
LOG.error("Cannot open websocket connection to Zeppelin note {}", noteId);
|
||||
return;
|
||||
}
|
||||
noteSession.getRemote().sendStringByFuture(serialize(msg));
|
||||
}
|
||||
|
||||
/* per notebook based ws connection, returns null if can't connect *
|
||||
public Session getZeppelinConnection(String noteId) {
|
||||
if (StringUtils.isBlank(noteId)) {
|
||||
LOG.warn("Cannot return websocket connection for blank noteId");
|
||||
LOG.warn("Cannot get Websocket session with blanck noteId");
|
||||
return null;
|
||||
}
|
||||
|
||||
if (zeppelinConnectionMap.containsKey(noteId)) {
|
||||
LOG.info("Connection for {} exists in map", noteId);
|
||||
return getNoteSession(noteId);
|
||||
}
|
||||
//TODO(khalid): clean log later
|
||||
LOG.info("Creating Zeppelin websocket connection {} {}", zeppelinWebsocketUrl, noteId);
|
||||
return openNoteSession(noteId);
|
||||
return getNoteSession(noteId);
|
||||
}
|
||||
*/
|
||||
|
||||
/*
|
||||
private Message zeppelinGetNoteMsg(String noteId) {
|
||||
Message getNoteMsg = new Message(Message.OP.GET_NOTE);
|
||||
HashMap<String, Object> data = new HashMap<>();
|
||||
|
|
@ -211,12 +206,14 @@ public class ZeppelinClient {
|
|||
getNoteMsg.data = data;
|
||||
return getNoteMsg;
|
||||
}
|
||||
/*
|
||||
*/
|
||||
|
||||
private Session getNoteSession(String noteId) {
|
||||
Session session = zeppelinConnectionMap.get(noteId);
|
||||
if (session == null || !session.isOpen()) {
|
||||
LOG.info("Not connection to {}", noteId);
|
||||
zeppelinConnectionMap.remove(noteId);
|
||||
LOG.info("Getting Note websocket connection for note {}", noteId);
|
||||
Session session = notesConnection.get(noteId);
|
||||
if (!isSessionOpen(session)) {
|
||||
LOG.info("No open connection for note {}, opening one", noteId);
|
||||
notesConnection.remove(noteId);
|
||||
session = openNoteSession(noteId);
|
||||
}
|
||||
return session;
|
||||
|
|
@ -235,18 +232,29 @@ public class ZeppelinClient {
|
|||
return session;
|
||||
}
|
||||
|
||||
if (zeppelinConnectionMap.containsKey(noteId)) {
|
||||
if (notesConnection.containsKey(noteId)) {
|
||||
session.close();
|
||||
session = zeppelinConnectionMap.get(noteId);
|
||||
session = notesConnection.get(noteId);
|
||||
} else {
|
||||
String getNote = serialize(zeppelinGetNoteMsg(noteId));
|
||||
// TODO(khalid): may need to check return whether successful
|
||||
session.getRemote().sendStringByFuture(getNote);
|
||||
zeppelinConnectionMap.put(noteId, session);
|
||||
notesConnection.put(noteId, session);
|
||||
}
|
||||
return session;
|
||||
}
|
||||
*/
|
||||
|
||||
private boolean isSessionOpen(Session session) {
|
||||
return (session != null) && (session.isOpen());
|
||||
}
|
||||
|
||||
private Message zeppelinGetNoteMsg(String noteId) {
|
||||
Message getNoteMsg = new Message(Message.OP.GET_NOTE);
|
||||
HashMap<String, Object> data = new HashMap<String, Object>();
|
||||
data.put("id", noteId);
|
||||
getNoteMsg.data = data;
|
||||
return getNoteMsg;
|
||||
}
|
||||
|
||||
public void handleMsgFromZeppelin(String message, String noteId) {
|
||||
Map<String, String> meta = new HashMap<>();
|
||||
meta.put("token", zeppelinhubToken);
|
||||
|
|
@ -264,51 +272,24 @@ public class ZeppelinClient {
|
|||
client.relayToZeppelinHub(hubMsg.serialize());
|
||||
}
|
||||
|
||||
/**
|
||||
* Close and remove ZeppelinConnection
|
||||
*
|
||||
public void removeZeppelinConnection(String noteId) {
|
||||
if (zeppelinConnectionMap.containsKey(noteId)) {
|
||||
Session conn = zeppelinConnectionMap.get(noteId);
|
||||
if (conn.isOpen()) {
|
||||
conn.close();
|
||||
}
|
||||
zeppelinConnectionMap.remove(noteId);
|
||||
public void removeNoteConnection(String noteId) {
|
||||
if (StringUtils.isBlank(noteId)) {
|
||||
LOG.error("Cannot remove session for empty noteId");
|
||||
return;
|
||||
}
|
||||
// TODO(khalid): clean log later
|
||||
LOG.info("Removed Zeppelin ws connection for the following note {}", noteId);
|
||||
if (notesConnection.containsKey(noteId)) {
|
||||
Session connection = notesConnection.get(noteId);
|
||||
if (connection.isOpen()) {
|
||||
connection.close();
|
||||
}
|
||||
notesConnection.remove(noteId);
|
||||
}
|
||||
LOG.info("Removed note websocket connection for note {}", noteId);
|
||||
}
|
||||
|
||||
/**
|
||||
* Close and remove all ZeppelinConnection
|
||||
*
|
||||
public void removeAllZeppelinConnections() {
|
||||
for (Map.Entry<String, Session> entry: zeppelinConnectionMap.entrySet()) {
|
||||
if (isSessionOpen(entry.getValue())) {
|
||||
entry.getValue().close();
|
||||
}
|
||||
zeppelinConnectionMap.remove(entry.getKey());
|
||||
}
|
||||
LOG.info("Removed all Zeppelin ws connections");
|
||||
}
|
||||
|
||||
public void pingAllNotes() {
|
||||
for (Map.Entry<String, Session> entry: zeppelinConnectionMap.entrySet()) {
|
||||
if (isSessionOpen(entry.getValue())) {
|
||||
send(new Message(OP.PING), entry.getKey());
|
||||
} else {
|
||||
// for cleanup
|
||||
zeppelinConnectionMap.remove(entry.getKey());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
public int countConnectedNotes() {
|
||||
return zeppelinConnectionMap.size();
|
||||
}
|
||||
*/
|
||||
public void ping() {
|
||||
if (watcherSession == null) {
|
||||
LOG.info("Cannot send PING event, watcher is null");
|
||||
return;
|
||||
}
|
||||
watcherSession.getRemote().sendStringByFuture(serialize(new Message(OP.PING)));
|
||||
|
|
|
|||
Loading…
Reference in a new issue