feat(broadcast): register docId -> broadcast server route when init connection with client

Author: BrotherJing

scalable-ot-broadcast/src/main/java/com/brotherjing/broadcast/handler/WebSocketHandler.java

@@ -12,13 +12,15 @@
 import lombok.extern.slf4j.Slf4j;
 
 import org.apache.dubbo.config.annotation.Service;
+import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.stereotype.Component;
 import org.springframework.util.CollectionUtils;
 import org.springframework.web.socket.BinaryMessage;
 import org.springframework.web.socket.WebSocketSession;
 import org.springframework.web.socket.handler.BinaryWebSocketHandler;
 
 import com.brotherjing.api.Broadcast;
+import com.brotherjing.broadcast.service.RegistryService;
 import com.brotherjing.core.util.Converter;
 import com.brotherjing.proto.BaseProto;
 
@@ -27,6 +29,9 @@
 @Service(interfaceClass = Broadcast.class)
 public class WebSocketHandler extends BinaryWebSocketHandler implements Broadcast {
 
+    @Autowired
+    private RegistryService registryService;
+
     private Map<String, Map<String, WebSocketSession>> clientsByDocId = new ConcurrentHashMap<>();
 
     @Override
@@ -104,5 +109,8 @@ public void sendTo(String sid, List<byte[]> commandsPayload) {
     private void register(String docId, String sid, WebSocketSession client) {
         clientsByDocId.putIfAbsent(docId, new ConcurrentHashMap<>());
         clientsByDocId.get(docId).put(sid, client);
+
+        // also register the route in redis
+        registryService.registerRoute(docId);
     }
 }

scalable-ot-broadcast/src/main/java/com/brotherjing/broadcast/service/RegistryService.java

@@ -1,23 +1,30 @@
 package com.brotherjing.broadcast.service;
 
+import java.net.InetAddress;
+import java.net.UnknownHostException;
+
 import javax.annotation.PostConstruct;
 
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.beans.factory.annotation.Value;
 import org.springframework.stereotype.Service;
 
-import com.brotherjing.Const;
 import com.brotherjing.config.ZookeeperConfig;
+import com.brotherjing.core.dao.RedisDao;
+import com.brotherjing.core.loadbalance.ServerEntity;
 import com.brotherjing.core.zookeeper.ZookeeperRegistry;
 
 @Service
 public class RegistryService {
 
     @Value(value = "${server.port}")
-    private String serverPort;
+    private int serverPort;
 
     @Value(value = "${dubbo.protocol.port}")
-    private String dubboPort;
+    private int dubboPort;
+
+    @Autowired
+    private RedisDao redisDao;
 
     @Autowired
     private ZookeeperConfig zookeeperConfig;
@@ -31,10 +38,38 @@ public void init() {
         register();
     }
 
+    /**
+     * Register this node in discovery service
+     */
     private void register() {
-        boolean success = registry.register(Const.BROADCAST_REGISTRY_PATH, serverPort, dubboPort);
+        boolean success = registry.register(getServerEntity());
         if (!success) {
             throw new RuntimeException("Failed to register");
         }
     }
+
+    /**
+     * Register a specific route from docId to node.
+     */
+    public void registerRoute(String docId) {
+        redisDao.addRoute(docId, getServerEntity());
+    }
+
+    private ServerEntity getServerEntity() {
+        return ServerEntity.builder()
+                           .host(getAddress())
+                           .serverPort(serverPort)
+                           .dubboPort(dubboPort)
+                           .build();
+    }
+
+    private String getAddress() {
+        InetAddress inetAddress;
+        try {
+            inetAddress = InetAddress.getLocalHost();
+        } catch (UnknownHostException e) {
+            throw new RuntimeException("Failed to get local address", e);
+        }
+        return inetAddress.getHostAddress();
+    }
 }

scalable-ot-core/src/main/java/com/brotherjing/core/zookeeper/ZookeeperRegistry.java

@@ -1,7 +1,5 @@
 package com.brotherjing.core.zookeeper;
 
-import java.net.InetAddress;
-import java.net.UnknownHostException;
 import java.util.concurrent.CountDownLatch;
 import java.util.function.BiConsumer;
 
@@ -15,7 +13,9 @@
 import org.apache.zookeeper.CreateMode;
 import org.apache.zookeeper.data.Stat;
 
+import com.brotherjing.Const;
 import com.brotherjing.config.ZookeeperConfig;
+import com.brotherjing.core.loadbalance.ServerEntity;
 
 @Slf4j
 public class ZookeeperRegistry {
@@ -79,21 +79,14 @@ public boolean discover(String path, BiConsumer<CuratorFramework, PathChildrenCa
      * Dubbo port is put in path, while server port is put in node data.
      * For example /root/127.0.0.1:20880 with data 127.0.0.1:8080
      */
-    public boolean register(String root, String serverPort, String dubboPort) {
-        InetAddress inetAddress;
-        try {
-            inetAddress = InetAddress.getLocalHost();
-        } catch (UnknownHostException e) {
-            throw new RuntimeException("Failed to get local address", e);
-        }
-        String address = inetAddress.getHostAddress();
-        String path = root + "/" + address + ":" + dubboPort;
-        String serverAddress = address + ":" + serverPort;
+    public boolean register(ServerEntity serverEntity) {
+        String path = Const.BROADCAST_REGISTRY_PATH + "/" + serverEntity.getDubboAddress();
+        String data = serverEntity.getServerAddress();
         try {
             curator.create()
                    .creatingParentsIfNeeded()
                    .withMode(CreateMode.EPHEMERAL)
-                   .forPath(path, serverAddress.getBytes());
+                   .forPath(path, data.getBytes());
             log.info("Registered path in zookeeper: {}", path);
         } catch (Exception e) {
             log.error("Failed to register path {}, {}", path, e);