From 29c2ae91ed28f2b1ac7e601669aac1214ca58cee Mon Sep 17 00:00:00 2001
From: egorklimov <klim.electronicmail@gmail.com>
Date: Wed, 29 Aug 2018 11:54:03 +0300
Subject: [PATCH] Exceptions added

---
 .../zeppelin/socket/ConnectionManager.java    | 20 +++++++++++--------
 1 file changed, 12 insertions(+), 8 deletions(-)

diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/socket/ConnectionManager.java b/zeppelin-server/src/main/java/org/apache/zeppelin/socket/ConnectionManager.java
index 5d02d9f53f3..ffadfa0ee55 100644
--- a/zeppelin-server/src/main/java/org/apache/zeppelin/socket/ConnectionManager.java
+++ b/zeppelin-server/src/main/java/org/apache/zeppelin/socket/ConnectionManager.java
@@ -34,6 +34,7 @@
 import org.apache.zeppelin.notebook.socket.WatcherMessage;
 import org.apache.zeppelin.user.AuthenticationInfo;
 import org.apache.zeppelin.util.WatcherSecurityKey;
+import org.eclipse.jetty.websocket.api.WebSocketException;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -213,7 +214,7 @@ public void broadcast(Message m) {
       for (NotebookSocket ns : connectedSockets) {
         try {
           ns.send(serializeMessage(m));
-        } catch (IOException e) {
+        } catch (IOException | WebSocketException e) {
           LOGGER.error("Send error: " + m, e);
         }
       }
@@ -234,7 +235,7 @@ public void broadcast(String noteId, Message m) {
     for (NotebookSocket conn : socketsToBroadcast) {
       try {
         conn.send(serializeMessage(m));
-      } catch (IOException e) {
+      } catch (IOException | WebSocketException e) {
         LOGGER.error("socket error", e);
       }
     }
@@ -245,9 +246,12 @@ private void broadcastToWatchers(String noteId, String subject, Message message)
       for (NotebookSocket watcher : watcherSockets) {
         try {
           watcher.send(
-              WatcherMessage.builder(noteId).subject(subject).message(serializeMessage(message))
-                  .build().toJson());
-        } catch (IOException e) {
+              WatcherMessage.builder(noteId)
+                  .subject(subject)
+                  .message(serializeMessage(message))
+                  .build()
+                  .toJson());
+        } catch (IOException | WebSocketException e) {
           LOGGER.error("Cannot broadcast message to watcher", e);
         }
       }
@@ -272,7 +276,7 @@ public void broadcastExcept(String noteId, Message m, NotebookSocket exclude) {
       }
       try {
         conn.send(serializeMessage(m));
-      } catch (IOException e) {
+      } catch (IOException | WebSocketException e) {
         LOGGER.error("socket error", e);
       }
     }
@@ -294,7 +298,7 @@ public void broadcastToAllConnectionsExcept(NotebookSocket exclude, String seria
 
         try {
           conn.send(serializedMsg);
-        } catch (IOException e) {
+        } catch (IOException | WebSocketException e) {
           LOGGER.error("Cannot broadcast message to conn", e);
         }
       }
@@ -324,7 +328,7 @@ public void multicastToUser(String user, Message m) {
   public void unicast(Message m, NotebookSocket conn) {
     try {
       conn.send(serializeMessage(m));
-    } catch (IOException e) {
+    } catch (IOException | WebSocketException e) {
       LOGGER.error("socket error", e);
     }
     broadcastToWatchers(StringUtils.EMPTY, StringUtils.EMPTY, m);