multi ws one way connection

zeppelin-server/src/main/java/org/apache/zeppelin/realm/ZeppelinHubRealm.java

@@ -38,6 +38,7 @@ import org.apache.shiro.authz.AuthorizationInfo;
 import org.apache.shiro.realm.AuthorizingRealm;
 import org.apache.shiro.subject.PrincipalCollection;
 import org.apache.zeppelin.notebook.repo.zeppelinhub.model.UserSessionContainer;
+import org.apache.zeppelin.notebook.repo.zeppelinhub.websocket.utils.ZeppelinhubUtils;
 import org.apache.zeppelin.server.ZeppelinServer;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -155,15 +156,8 @@ public class ZeppelinHubRealm extends AuthorizingRealm {
       throw new AuthenticationException("Cannot login to ZeppelinHub");
     }
 
-    // Add ZeppelinHub user_session token this singleton map, this will help ZeppelinHubRepo
-    // to get specific information about the current user.
-    UserSessionContainer.instance.setSession(account.login, userSession);
-
-    /* TODO(khalid): add proper roles and add listener */
-    HashSet<String> userAndRoles = new HashSet<String>();
-    userAndRoles.add(account.login);
-    ZeppelinServer.notebookWsServer.broadcastReloadedNoteList(
-        new org.apache.zeppelin.user.AuthenticationInfo(account.login), userAndRoles);
+    onLoginSuccess(account.login, userSession);
+    
     return account;
   }
 
@@ -213,4 +207,22 @@ public class ZeppelinHubRealm extends AuthorizingRealm {
     public String email;
     public String name;
   }
+  
+  public void onLoginSuccess(String username, String session) {
+    UserSessionContainer.instance.setSession(username, session);
+
+    /* TODO(xxx): add proper roles */
+    HashSet<String> userAndRoles = new HashSet<String>();
+    userAndRoles.add(username);
+    ZeppelinServer.notebookWsServer.broadcastReloadedNoteList(
+        new org.apache.zeppelin.user.AuthenticationInfo(username), userAndRoles);
+
+    LOG.info("broadcasted notes, starting userLoginRoutine");
+    ZeppelinhubUtils.userLoginRoutine(username);
+  }
+  
+  @Override
+  public void onLogout(PrincipalCollection principals) {
+    ZeppelinhubUtils.userLogoutRoutine((String) principals.getPrimaryPrincipal());
+  }
 }

zeppelin-zengine/src/main/java/org/apache/zeppelin/notebook/repo/zeppelinhub/ZeppelinHubRepo.java

@@ -22,8 +22,6 @@ import java.net.URISyntaxException;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
-import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.ConcurrentMap;
 
 import org.apache.commons.lang.StringUtils;
 import org.apache.zeppelin.conf.ZeppelinConfiguration;
@@ -32,8 +30,10 @@ import org.apache.zeppelin.notebook.NoteInfo;
 import org.apache.zeppelin.notebook.repo.NotebookRepo;
 import org.apache.zeppelin.notebook.repo.NotebookRepoSettingsInfo;
 import org.apache.zeppelin.notebook.repo.zeppelinhub.model.Instance;
+import org.apache.zeppelin.notebook.repo.zeppelinhub.model.UserTokenContainer;
 import org.apache.zeppelin.notebook.repo.zeppelinhub.model.UserSessionContainer;
 import org.apache.zeppelin.notebook.repo.zeppelinhub.rest.ZeppelinhubRestApiHandler;
+import org.apache.zeppelin.notebook.repo.zeppelinhub.websocket.Client;
 import org.apache.zeppelin.user.AuthenticationInfo;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -55,27 +55,28 @@ public class ZeppelinHubRepo implements NotebookRepo {
   public static final String TOKEN_HEADER = "X-Zeppelin-Token";
   private static final Gson GSON = new Gson();
   private static final Note EMPTY_NOTE = new Note();
-  //private final Client websocketClient;
+  private final Client websocketClient;
+  private final UserTokenContainer tokenManager;
 
   private String token;
   private ZeppelinhubRestApiHandler restApiClient;
   
   private final ZeppelinConfiguration conf;
-
-  // In order to avoid too many call to ZeppelinHub backend, we save a map of user -> session.
-  private ConcurrentMap<String, String> usersToken = new ConcurrentHashMap<String, String>();
   
   public ZeppelinHubRepo(ZeppelinConfiguration conf) {
     this.conf = conf;
     String zeppelinHubUrl = getZeppelinHubUrl(conf);
     LOG.info("Initializing ZeppelinHub integration module");
-    token = conf.getString("ZEPPELINHUB_API_TOKEN", ZEPPELIN_CONF_PROP_NAME_TOKEN, "");
-    restApiClient = ZeppelinhubRestApiHandler.newInstance(zeppelinHubUrl, token);
 
-    // TODO(xxx): refactor this in the next itaration
-    //websocketClient = Client.initialize(getZeppelinWebsocketUri(conf),
-    //    getZeppelinhubWebsocketUri(conf), token, conf);
-    //websocketClient.start();
+    token = conf.getString("ZEPPELINHUB_API_TOKEN", ZEPPELIN_CONF_PROP_NAME_TOKEN, "");
+    restApiClient = ZeppelinhubRestApiHandler.newInstance(zeppelinHubUrl);
+    //TODO(khalid): check which realm for authentication, pass to token manager
+    tokenManager = UserTokenContainer.init(restApiClient, token);
+
+    websocketClient = Client.initialize(getZeppelinWebsocketUri(conf),
+        getZeppelinhubWebsocketUri(conf), token, conf);
+    websocketClient.start();
+
   }
 
   private String getZeppelinHubWsUri(URI api) throws URISyntaxException {
@@ -155,58 +156,6 @@ public class ZeppelinHubRepo implements NotebookRepo {
     }
     return zeppelinhubUrl;
   }
-  
-  /**
-   * Get list of user instances from Zeppelinhub.
-   * This will avoid and remove the needs of setting up token in zeppelin-env.sh.
-   */
-  private List<Instance> getUserInstances(String ticket) throws IOException {
-    if (StringUtils.isBlank(ticket)) {
-      return Collections.emptyList();
-    }
-    return restApiClient.getInstances(ticket);
-  }
-
-  /**
-   * Get user default instance.
-   * From now, it will be from the first instance from the list,
-   * But later we can think about marking a default one and return it instead :)
-   */
-  private String getDefaultZeppelinInstanceToken(String ticket) throws IOException {
-    List<Instance> instances = getUserInstances(ticket);
-    if (instances.isEmpty()) {
-      return StringUtils.EMPTY;
-    }
-
-    String token = instances.get(0).token;
-    LOG.debug("The following instance has been assigned {} with token {}", instances.get(0).name,
-        token);
-    return token;
-  }
-
-  /**
-   * For a given user logged in is zeppelin (via zeppelinhub notebook repo), get default token.
-   *
-   */
-  private String getUserToken(String principal) {
-    // Case of user use token instead of authentication.
-    if (!StringUtils.isBlank(token)) {
-      return token;
-    }
-
-    String token = usersToken.get(principal);
-    if (StringUtils.isBlank(token)) {
-      String ticket = UserSessionContainer.instance.getSession(principal);
-      try {
-        token = getDefaultZeppelinInstanceToken(ticket);
-        usersToken.putIfAbsent(principal, token);
-      } catch (IOException e) {
-        LOG.error("Cannot get user token", e);
-        token = StringUtils.EMPTY;
-      }
-    }
-    return token;
-  }
 
   private boolean isSubjectValid(AuthenticationInfo subject) {
     if (subject == null) {
@@ -320,6 +269,10 @@ public class ZeppelinHubRepo implements NotebookRepo {
     }
     return history;
   }
+  
+  private String getUserToken(String user) {
+    return tokenManager.getUserToken(user);
+  }
 
   @Override
   public List<NotebookRepoSettingsInfo> getSettings(AuthenticationInfo subject) {
@@ -335,7 +288,7 @@ public class ZeppelinHubRepo implements NotebookRepo {
     List<Map<String, String>> values = Lists.newLinkedList();
 
     try {
-      instances = getUserInstances(zeppelinHubUserSession);
+      instances = tokenManager.getUserInstances(zeppelinHubUserSession);
     } catch (IOException e) {
       LOG.warn("Couldnt find instances for the session {}, returning empty collection",
           zeppelinHubUserSession);
@@ -369,7 +322,7 @@ public class ZeppelinHubRepo implements NotebookRepo {
     String ticket = UserSessionContainer.instance.getSession(user);
     List<Instance> instances;
     try {
-      instances = getUserInstances(ticket);
+      instances = tokenManager.getUserInstances(ticket);
       if (instances.isEmpty()) {
         return;
       }
@@ -377,7 +330,7 @@ public class ZeppelinHubRepo implements NotebookRepo {
       for (Instance instance : instances) {
         if (instance.id == instanceId) {
           LOG.info("User {} switched to instance {}", user, instances.get(0).name);
-          usersToken.put(user, instance.token);
+          tokenManager.setUserToken(user, instance.token);
           break;
         }
       }

zeppelin-zengine/src/main/java/org/apache/zeppelin/notebook/repo/zeppelinhub/model/UserSessionContainer.java

@@ -28,7 +28,7 @@ import org.apache.commons.lang.StringUtils;
 public class UserSessionContainer {
   private static class Entity {
     public final String userSession;
-
+    
     Entity(String userSession) {
       this.userSession = userSession;
     }

zeppelin-zengine/src/main/java/org/apache/zeppelin/notebook/repo/zeppelinhub/model/UserTokenContainer.java

@@ -0,0 +1,122 @@
+/*
+ * 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.notebook.repo.zeppelinhub.model;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
+
+import org.apache.commons.lang.StringUtils;
+import org.apache.zeppelin.notebook.repo.zeppelinhub.rest.ZeppelinhubRestApiHandler;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * User token manager class.
+ *
+ */
+
+public class UserTokenContainer {
+  private static final Logger LOG = LoggerFactory.getLogger(UserTokenContainer.class);
+  public static UserTokenContainer instance = null;
+  private ConcurrentMap<String, String> userTokens = new ConcurrentHashMap<String, String>();
+  private final ZeppelinhubRestApiHandler restApiClient;
+  private String defaultToken;
+
+  public static UserTokenContainer init(ZeppelinhubRestApiHandler restClient, 
+      String defaultToken) {
+    if (instance == null) {
+      instance = new UserTokenContainer(restClient, defaultToken);
+    }
+    return instance;
+  }
+
+  private UserTokenContainer(ZeppelinhubRestApiHandler restClient, String defaultToken) {
+    LOG.info("Initializing ZeppelinHub token container");
+    restApiClient = restClient;
+    this.defaultToken = defaultToken;
+  }
+
+  public void setUserToken(String username, String token) {
+    userTokens.put(username, token);
+  }
+  
+  public String getUserToken(String principal) {
+    // Case of user use token instead of authentication.
+    //TODO(khalid): in case of hub realm, remote token should take precedence over default
+    if (!StringUtils.isBlank(defaultToken)) {
+      return defaultToken;
+    }
+    String token = userTokens.get(principal);
+    if (StringUtils.isBlank(token)) {
+      String ticket = UserSessionContainer.instance.getSession(principal);
+      try {
+        token = getDefaultZeppelinInstanceToken(ticket);
+        userTokens.putIfAbsent(principal, token);
+      } catch (IOException e) {
+        LOG.error("Cannot get user token", e);
+        token = StringUtils.EMPTY;
+      }
+    }
+    return token;
+  }
+  
+  public String removeUserToken(String username) {
+    return userTokens.remove(username);
+  }
+  
+  /**
+   * Get user default instance.
+   * From now, it will be from the first instance from the list,
+   * But later we can think about marking a default one and return it instead :)
+   */
+  public String getDefaultZeppelinInstanceToken(String ticket) throws IOException {
+    List<Instance> instances = getUserInstances(ticket);
+    if (instances.isEmpty()) {
+      return StringUtils.EMPTY;
+    }
+
+    String token = instances.get(0).token;
+    LOG.debug("The following instance has been assigned {} with token {}", instances.get(0).name,
+        token);
+    return token;
+  }
+  
+  /**
+   * Get list of user instances from Zeppelinhub.
+   * This will avoid and remove the needs of setting up token in zeppelin-env.sh.
+   */
+  public List<Instance> getUserInstances(String ticket) throws IOException {
+    if (StringUtils.isBlank(ticket)) {
+      return Collections.emptyList();
+    }
+    return restApiClient.getInstances(ticket);
+  }
+  
+  public List<String> getAllTokens() {
+    return new ArrayList<String>(userTokens.values());
+  }
+  
+  public Map<String, String> getAllUserTokens() {
+    return new HashMap<String, String>(userTokens);
+  }
+}

zeppelin-zengine/src/main/java/org/apache/zeppelin/notebook/repo/zeppelinhub/rest/ZeppelinhubRestApiHandler.java

@@ -58,11 +58,11 @@ public class ZeppelinhubRestApiHandler {
   private final HttpClient client;
   private final String zepelinhubUrl;
 
-  public static ZeppelinhubRestApiHandler newInstance(String zeppelinhubUrl, String token) {
-    return new ZeppelinhubRestApiHandler(zeppelinhubUrl, token);
+  public static ZeppelinhubRestApiHandler newInstance(String zeppelinhubUrl) {
+    return new ZeppelinhubRestApiHandler(zeppelinhubUrl);
   }
 
-  private ZeppelinhubRestApiHandler(String zeppelinhubUrl, String token) {
+  private ZeppelinhubRestApiHandler(String zeppelinhubUrl) {
     this.zepelinhubUrl = zeppelinhubUrl + DEFAULT_API_PATH + "/";
 
     //TODO(khalid):to make proxy conf consistent with Zeppelin confs

zeppelin-zengine/src/main/java/org/apache/zeppelin/notebook/repo/zeppelinhub/websocket/Client.java

@@ -72,8 +72,8 @@ public class Client {
     }
   }
 
-  public void relayToZeppelinHub(String message) {
-    zeppelinhubClient.send(message);
+  public void relayToZeppelinHub(String message, String token) {
+    zeppelinhubClient.send(message, token);
   }
 
   public void relayToZeppelin(Message message, String noteId) {

zeppelin-zengine/src/main/java/org/apache/zeppelin/notebook/repo/zeppelinhub/websocket/ZeppelinClient.java

@@ -18,8 +18,11 @@ package org.apache.zeppelin.notebook.repo.zeppelinhub.websocket;
 
 import java.io.IOException;
 import java.net.URI;
+import java.util.Arrays;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.Map;
+import java.util.Set;
 import java.util.Timer;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ExecutionException;
@@ -27,6 +30,8 @@ import java.util.concurrent.Future;
 
 import org.apache.commons.lang3.StringUtils;
 import org.apache.zeppelin.conf.ZeppelinConfiguration;
+import org.apache.zeppelin.notebook.NotebookAuthorization;
+import org.apache.zeppelin.notebook.repo.zeppelinhub.model.UserTokenContainer;
 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;
@@ -53,7 +58,6 @@ import com.google.gson.JsonSyntaxException;
 public class ZeppelinClient {
   private static final Logger LOG = LoggerFactory.getLogger(ZeppelinClient.class);
   private final URI zeppelinWebsocketUrl;
-  private final String zeppelinhubToken;
   private final WebSocketClient wsClient;
   private static Gson gson;
   // Keep track of current open connection per notebook.
@@ -65,6 +69,12 @@ public class ZeppelinClient {
   private Authentication authModule;
   private static final int MIN = 60;
 
+  private static final Set<String> actionable = new  HashSet<String>(Arrays.asList(
+      "PROGRESS",
+      "NOTE",
+      "PARAGRAPH",
+      "PARAGRAPH_UPDATE_OUTPUT"));
+
   public static ZeppelinClient initialize(String zeppelinUrl, String token, 
       ZeppelinConfiguration conf) {
     if (instance == null) {
@@ -79,7 +89,6 @@ public class ZeppelinClient {
 
   private ZeppelinClient(String zeppelinUrl, String token, ZeppelinConfiguration conf) {
     zeppelinWebsocketUrl = URI.create(zeppelinUrl);
-    zeppelinhubToken = token;
     wsClient = createNewWebsocketClient();
     gson = new Gson();
     notesConnection = new ConcurrentHashMap<>();
@@ -121,7 +130,7 @@ public class ZeppelinClient {
       public void run() {
         watcherSession = openWatcherSession();
       }
-    }, 5000);
+    }, 10000);
   }
 
   public void stop() {
@@ -260,21 +269,64 @@ public class ZeppelinClient {
 
   public void handleMsgFromZeppelin(String message, String noteId) {
     Map<String, String> meta = new HashMap<>();
-    meta.put("token", zeppelinhubToken);
+    //TODO(khalid): don't use zeppelinhubToken in this class, decouple
     meta.put("noteId", noteId);
     Message zeppelinMsg = deserialize(message);
     if (zeppelinMsg == null) {
       return;
     }
-    ZeppelinhubMessage hubMsg = ZeppelinhubMessage.newMessage(zeppelinMsg, meta);
+    LOG.info("message from {}", zeppelinMsg.principal);
+    String token;
+    if (!isActionable(zeppelinMsg.op)) {
+      return;
+    }
+    if (StringUtils.isEmpty(zeppelinMsg.principal) || zeppelinMsg.principal == "anonymous") {
+      token = UserTokenContainer.instance.getUserToken(zeppelinMsg.principal);
+    } else {
+      token = "anonymous";
+    }
+    meta.put("token", token);
     Client client = Client.getInstance();
     if (client == null) {
       LOG.warn("Client isn't initialized yet");
       return;
     }
-    client.relayToZeppelinHub(hubMsg.serialize());
+    ZeppelinhubMessage hubMsg = ZeppelinhubMessage.newMessage(zeppelinMsg, meta);
+    if (token == "anonymous") {
+      relayToAllZeppelinHub(hubMsg, noteId);
+    } else {
+      client.relayToZeppelinHub(hubMsg.serialize(), zeppelinMsg.ticket);
+    }
+
   }
 
+  private void relayToAllZeppelinHub(ZeppelinhubMessage hubMsg, String noteId) {
+    if (StringUtils.isBlank(noteId)) {
+      return;
+    }
+    NotebookAuthorization noteAuth = NotebookAuthorization.getInstance();
+    Map<String, String> userTokens = UserTokenContainer.instance.getAllUserTokens();
+    Client client = Client.getInstance();
+    Set<String> userAndRoles;
+    String token;
+    for (String user: userTokens.keySet()) {
+      userAndRoles = noteAuth.getRoles(user);
+      userAndRoles.add(user);
+      if (noteAuth.isReader(noteId, userAndRoles)) {
+        token = userTokens.get(user);
+        hubMsg.meta.put("token", token);
+        client.relayToZeppelinHub(hubMsg.serialize(), token);
+      }
+    }
+  }
+
+  private boolean isActionable(OP action) {
+    if (action == null) {
+      return false;
+    }
+    return actionable.contains(action.name());
+  }
+  
   public void removeNoteConnection(String noteId) {
     if (StringUtils.isBlank(noteId)) {
       LOG.error("Cannot remove session for empty noteId");
@@ -307,7 +359,7 @@ public class ZeppelinClient {
 
   public void ping() {
     if (watcherSession == null) {
-      LOG.info("Cannot send PING event, no watcher found");
+      LOG.debug("Cannot send PING event, no watcher found");
       return;
     }
     watcherSession.getRemote().sendStringByFuture(serialize(new Message(OP.PING)));

zeppelin-zengine/src/main/java/org/apache/zeppelin/notebook/repo/zeppelinhub/websocket/ZeppelinhubClient.java

@@ -21,6 +21,7 @@ import java.io.IOException;
 import java.net.HttpCookie;
 import java.net.URI;
 import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
@@ -58,7 +59,6 @@ public class ZeppelinhubClient {
 
   private final WebSocketClient client;
   private final URI zeppelinhubWebsocketUrl;
-  private final ClientUpgradeRequest conectionRequest;
   private final String zeppelinhubToken;
 
   private static final long CONNECTION_IDLE_TIME = TimeUnit.SECONDS.toMillis(30);
@@ -66,7 +66,8 @@ public class ZeppelinhubClient {
   private static Gson gson;
   
   private SchedulerService schedulerService;
-  private ZeppelinhubSession zeppelinhubSession;
+  private Map<String, ZeppelinhubSession> sessionMap = 
+      new ConcurrentHashMap<String, ZeppelinhubSession>();
 
   public static ZeppelinhubClient initialize(String zeppelinhubUrl, String token) {
     if (instance == null) {
@@ -82,7 +83,6 @@ public class ZeppelinhubClient {
   private ZeppelinhubClient(String url, String token) {
     zeppelinhubWebsocketUrl = URI.create(url);
     client = createNewWebsocketClient();
-    conectionRequest = setConnectionrequest(token);
     zeppelinhubToken = token;
     schedulerService = SchedulerService.create(10);
     gson = new Gson();
@@ -92,17 +92,19 @@ public class ZeppelinhubClient {
   public void start() {
     try {
       client.start();
-      zeppelinhubSession = connect();
       addRoutines();
     } catch (Exception e) {
       LOG.error("Cannot connect to zeppelinhub via websocket", e);
     }
   }
   
+  public void initUser(String token) {
+    
+  }
+
   public void stop() {
     LOG.info("Stopping Zeppelinhub websocket client");
     try {
-      zeppelinhubSession.close();
       schedulerService.close();
       client.stop();
     } catch (Exception e) {
@@ -110,14 +112,19 @@ public class ZeppelinhubClient {
     }
   }
 
+  public void stopUser(String token) {
+    removeSession(token);
+  }
+
   public String getToken() {
     return this.zeppelinhubToken;
   }
   
-  public void send(String msg) {
-    if (!isConnectedToZeppelinhub()) {
+  public void send(String msg, String token) {
+    ZeppelinhubSession zeppelinhubSession = getSession(token);
+    if (!isConnectedToZeppelinhub(zeppelinhubSession)) {
       LOG.info("Zeppelinhub connection is not open, opening it");
-      zeppelinhubSession = connect();
+      zeppelinhubSession = connect(token);
       if (zeppelinhubSession == ZeppelinhubSession.EMPTY) {
         LOG.warn("While connecting to ZeppelinHub received empty session, cannot send the message");
         return;
@@ -126,25 +133,48 @@ public class ZeppelinhubClient {
     zeppelinhubSession.sendByFuture(msg);
   }
   
-  private boolean isConnectedToZeppelinhub() {
+  private boolean isConnectedToZeppelinhub(ZeppelinhubSession zeppelinhubSession) {
     return (zeppelinhubSession != null && zeppelinhubSession.isSessionOpen());
   }
 
-  private ZeppelinhubSession connect() {
-    ZeppelinhubSession zeppelinSession;
-    try {
-      ZeppelinhubWebsocket ws = ZeppelinhubWebsocket.newInstance(zeppelinhubToken);
-      Future<Session> future = client.connect(ws, zeppelinhubWebsocketUrl, conectionRequest);
-      Session session = future.get();
-      zeppelinSession = ZeppelinhubSession.createInstance(session, zeppelinhubToken);
-    } catch (IOException | InterruptedException | ExecutionException e) {
-      LOG.info("Couldnt connect to zeppelinhub - {}", e.toString());
-      zeppelinSession = ZeppelinhubSession.EMPTY;
+  private ZeppelinhubSession connect(String token) {
+    if (StringUtils.isBlank(token)) {
+      LOG.debug("Can't connect with empty token");
+      return ZeppelinhubSession.EMPTY;
     }
-    return zeppelinSession;
+    ZeppelinhubSession zeppelinhubSession;
+    try {
+      ZeppelinhubWebsocket ws = ZeppelinhubWebsocket.newInstance(token);
+      ClientUpgradeRequest request = getConnectionRequest(token);
+      Future<Session> future = client.connect(ws, zeppelinhubWebsocketUrl, request);
+      Session session = future.get();
+      zeppelinhubSession = ZeppelinhubSession.createInstance(session, token);
+      setSession(token, zeppelinhubSession);
+    } catch (IOException | InterruptedException | ExecutionException e) {
+      LOG.info("Couldnt connect to zeppelinhub", e);
+      zeppelinhubSession = ZeppelinhubSession.EMPTY;
+    }
+    return zeppelinhubSession;
   }
-  
-  private ClientUpgradeRequest setConnectionrequest(String token) {
+
+  private void setSession(String token, ZeppelinhubSession session) {
+    sessionMap.put(token, session);
+  }
+
+  private ZeppelinhubSession getSession(String token) {
+    return sessionMap.get(token);
+  }
+
+  public void removeSession(String token) {
+    ZeppelinhubSession zeppelinhubSession = getSession(token);
+    if (zeppelinhubSession == null) {
+      return;
+    }
+    zeppelinhubSession.close();
+    sessionMap.remove(token);
+  }
+
+  private ClientUpgradeRequest getConnectionRequest(String token) {
     ClientUpgradeRequest request = new ClientUpgradeRequest();
     request.setCookies(Lists.newArrayList(new HttpCookie(ZeppelinHubRepo.TOKEN_HEADER, token)));
     return request;
@@ -193,7 +223,7 @@ public class ZeppelinhubClient {
           runAllParagraph(hubMsg.meta.get("noteId"), msg);
           break;
         default:
-          LOG.warn("Received {} from ZeppelinHub, not handled", op);
+          LOG.debug("Received {} from ZeppelinHub, not handled", op);
           break;
     }
   }

zeppelin-zengine/src/main/java/org/apache/zeppelin/notebook/repo/zeppelinhub/websocket/listener/ZeppelinhubWebsocket.java

@@ -53,7 +53,7 @@ public class ZeppelinhubWebsocket implements WebSocketListener {
   public void onWebSocketConnect(Session session) {
     LOG.info("Opening a new session to Zeppelinhub {}", session.hashCode());
     this.zeppelinHubSession = session;
-    send(ZeppelinhubUtils.liveMessage(token));
+    //send(ZeppelinhubUtils.liveMessage(token));
   }
 
   @Override

zeppelin-zengine/src/main/java/org/apache/zeppelin/notebook/repo/zeppelinhub/websocket/scheduler/ZeppelinHubHeartbeat.java

@@ -16,6 +16,7 @@
  */
 package org.apache.zeppelin.notebook.repo.zeppelinhub.websocket.scheduler;
 
+import org.apache.zeppelin.notebook.repo.zeppelinhub.model.UserTokenContainer;
 import org.apache.zeppelin.notebook.repo.zeppelinhub.websocket.ZeppelinhubClient;
 import org.apache.zeppelin.notebook.repo.zeppelinhub.websocket.utils.ZeppelinhubUtils;
 import org.slf4j.Logger;
@@ -40,6 +41,9 @@ public class ZeppelinHubHeartbeat implements Runnable {
   @Override
   public void run() {
     LOG.debug("Sending PING to zeppelinhub");
-    client.send(ZeppelinhubUtils.pingMessage(client.getToken()));
+    for (String token: UserTokenContainer.instance.getAllTokens()) {
+      LOG.info("Sending PING to zeppelinhub token {}", token);
+      client.send(ZeppelinhubUtils.pingMessage(token), token);
+    }
   }  
 }

zeppelin-zengine/src/main/java/org/apache/zeppelin/notebook/repo/zeppelinhub/websocket/utils/ZeppelinhubUtils.java

@@ -19,6 +19,8 @@ package org.apache.zeppelin.notebook.repo.zeppelinhub.websocket.utils;
 import java.util.HashMap;
 
 import org.apache.commons.lang.StringUtils;
+import org.apache.zeppelin.notebook.repo.zeppelinhub.model.UserTokenContainer;
+import org.apache.zeppelin.notebook.repo.zeppelinhub.websocket.ZeppelinhubClient;
 import org.apache.zeppelin.notebook.repo.zeppelinhub.websocket.protocol.ZeppelinHubOp;
 import org.apache.zeppelin.notebook.repo.zeppelinhub.websocket.protocol.ZeppelinhubMessage;
 import org.apache.zeppelin.notebook.socket.Message;
@@ -95,4 +97,25 @@ public class ZeppelinhubUtils {
   public static boolean isZeppelinOp(String text) {
     return (toZeppelinOp(text) != null); 
   }
+  
+  public static void userLoginRoutine(String username) {
+    LOG.info("Executing user login routine");
+    String token = UserTokenContainer.instance.getUserToken(username);
+    UserTokenContainer.instance.setUserToken(username, token);
+    LOG.info("Token for {} is {} ", username, token);
+    String msg = ZeppelinhubUtils.liveMessage(token);
+    LOG.info("Live message is {} with token {}", msg, token);
+    ZeppelinhubClient.getInstance()
+        .send(msg, token);
+  }
+  
+  public static void userLogoutRoutine(String username) {
+    LOG.info("Executing user logout routine");
+    String token = UserTokenContainer.instance.removeUserToken(username);
+    String msg = ZeppelinhubUtils.deadMessage(token);
+    LOG.info("Dead message is {} with token {}", msg, token);
+    ZeppelinhubClient.getInstance()
+        .send(msg, token);
+    ZeppelinhubClient.getInstance().removeSession(token);
+  }
 }