Handle ServerUnreachableException and and ServerNotFoundException as NO_SUB for transient session (#120)

THIS IS A TRADE-OFF! 
Reason: A network partition (which should be rare compared to a server crash) may also cause SUE and SNFE, leading to unexpected clearing of route information. Therefore, manual intervention is required if this happens.

Author: Gujiawei-Edinburgh

bifromq-mqtt/bifromq-mqtt-broker-client/pom.xml

@@ -41,5 +41,10 @@
             <groupId>org.slf4j</groupId>
             <artifactId>slf4j-api</artifactId>
         </dependency>
+        <dependency>
+            <groupId>org.mockito</groupId>
+            <artifactId>mockito-core</artifactId>
+            <scope>test</scope>
+        </dependency>
     </dependencies>
 </project>

bifromq-mqtt/bifromq-mqtt-broker-client/src/main/java/com/baidu/bifromq/mqtt/inbox/MqttBrokerClient.java

@@ -17,6 +17,8 @@
 import static java.util.Collections.emptyMap;
 
 import com.baidu.bifromq.baserpc.IRPCClient;
+import com.baidu.bifromq.baserpc.exception.ServerNotFoundException;
+import com.baidu.bifromq.baserpc.exception.ServerUnreachableException;
 import com.baidu.bifromq.mqtt.inbox.rpc.proto.OnlineInboxBrokerGrpc;
 import com.baidu.bifromq.mqtt.inbox.rpc.proto.SubReply;
 import com.baidu.bifromq.mqtt.inbox.rpc.proto.SubRequest;
@@ -25,12 +27,13 @@
 import com.baidu.bifromq.mqtt.inbox.rpc.proto.WriteReply;
 import com.baidu.bifromq.mqtt.inbox.rpc.proto.WriteRequest;
 import com.baidu.bifromq.mqtt.inbox.util.DeliveryGroupKeyUtil;
-import com.baidu.bifromq.plugin.subbroker.DeliveryReply;
-import com.baidu.bifromq.plugin.subbroker.DeliveryRequest;
-import com.baidu.bifromq.plugin.subbroker.IDeliverer;
+import com.baidu.bifromq.plugin.subbroker.*;
+import com.baidu.bifromq.type.MatchInfo;
 import com.baidu.bifromq.type.QoS;
 import com.google.common.base.Preconditions;
 import io.reactivex.rxjava3.core.Observable;
+
+import java.util.*;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.atomic.AtomicBoolean;
 import lombok.extern.slf4j.Slf4j;
@@ -40,14 +43,8 @@ final class MqttBrokerClient implements IMqttBrokerClient {
     private final AtomicBoolean hasStopped = new AtomicBoolean();
     private final IRPCClient rpcClient;
 
-    MqttBrokerClient(MqttBrokerClientBuilder builder) {
-        this.rpcClient = IRPCClient.newBuilder()
-            .bluePrint(RPCBluePrint.INSTANCE)
-            .executor(builder.executor)
-            .eventLoopGroup(builder.eventLoopGroup)
-            .sslContext(builder.sslContext)
-            .crdtService(builder.crdtService)
-            .build();
+    MqttBrokerClient(IRPCClient rpcClient) {
+        this.rpcClient = rpcClient;
     }
 
     public IDeliverer open(String delivererKey) {
@@ -127,7 +124,33 @@ public CompletableFuture<DeliveryReply> deliver(DeliveryRequest request) {
                     .setReqId(reqId)
                     .setRequest(request)
                     .build())
-                .thenApply(WriteReply::getReply);
+                .thenApply(WriteReply::getReply)
+                .handle((reply, ex) -> {
+                    if (ex != null) {
+                        if (ex.getCause() instanceof ServerUnreachableException
+                                || ex.getCause() instanceof ServerNotFoundException) {
+                            DeliveryReply.Builder replyBuilder = DeliveryReply.newBuilder();
+                            Map<String, DeliveryResults> resultsMap = new HashMap<>();
+                            for (Map.Entry<String, DeliveryPackage> entry : request.getPackageMap().entrySet()) {
+                                String tenantId = entry.getKey();
+                                Set<MatchInfo> noSub = new HashSet<>();
+                                DeliveryResults.Builder deliveryResultsBuilder = DeliveryResults.newBuilder();
+                                for (DeliveryPack pack: entry.getValue().getPackList()) {
+                                    noSub.addAll(pack.getMatchInfoList());
+                                }
+                                noSub.forEach(matchInfo -> deliveryResultsBuilder.addResult(DeliveryResult.newBuilder()
+                                        .setMatchInfo(matchInfo)
+                                        .setCode(DeliveryResult.Code.NO_SUB)
+                                        .build()));
+                                resultsMap.put(tenantId, deliveryResultsBuilder.build());
+                            }
+                            return  replyBuilder.putAllResult(resultsMap).build();
+                        } else {
+                            throw new RuntimeException(ex);
+                        }
+                    }
+                    return reply;
+                });
         }
 
         @Override

bifromq-mqtt/bifromq-mqtt-broker-client/src/main/java/com/baidu/bifromq/mqtt/inbox/MqttBrokerClientBuilder.java

@@ -14,6 +14,7 @@
 package com.baidu.bifromq.mqtt.inbox;
 
 import com.baidu.bifromq.basecrdt.service.ICRDTService;
+import com.baidu.bifromq.baserpc.IRPCClient;
 import io.netty.channel.EventLoopGroup;
 import io.netty.handler.ssl.SslContext;
 import java.util.concurrent.Executor;
@@ -32,6 +33,12 @@ public final class MqttBrokerClientBuilder implements IMqttBrokerClientBuilder {
     Executor executor;
 
     public IMqttBrokerClient build() {
-        return new MqttBrokerClient(this);
+        return new MqttBrokerClient(IRPCClient.newBuilder()
+                .bluePrint(RPCBluePrint.INSTANCE)
+                .executor(executor)
+                .eventLoopGroup(eventLoopGroup)
+                .sslContext(sslContext)
+                .crdtService(crdtService)
+                .build());
     }
 }

bifromq-mqtt/bifromq-mqtt-broker-client/src/test/java/com/baidu/bifromq/mqtt/inbox/MqttBrokerClientTest.java

@@ -0,0 +1,102 @@
+package com.baidu.bifromq.mqtt.inbox;
+
+import com.baidu.bifromq.baserpc.IRPCClient;
+import com.baidu.bifromq.baserpc.exception.ServerNotFoundException;
+import com.baidu.bifromq.mqtt.inbox.rpc.proto.OnlineInboxBrokerGrpc;
+import com.baidu.bifromq.mqtt.inbox.rpc.proto.WriteReply;
+import com.baidu.bifromq.mqtt.inbox.rpc.proto.WriteRequest;
+import com.baidu.bifromq.mqtt.inbox.util.DeliveryGroupKeyUtil;
+import com.baidu.bifromq.plugin.subbroker.*;
+import com.baidu.bifromq.type.MatchInfo;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+import org.testng.annotations.AfterMethod;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+
+import java.util.concurrent.CompletableFuture;
+
+import static java.util.Collections.emptyMap;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.when;
+
+public class MqttBrokerClientTest {
+    private AutoCloseable closeable;
+    @Mock
+    private IRPCClient rpcClient;
+    private MqttBrokerClient mqttBrokerClient;
+    @Mock
+    private IDeliverer deliverer;
+    @Mock
+    private IRPCClient.IRequestPipeline<WriteRequest, WriteReply> pipeline;
+    private final String delivererKey = "test:DelivererKey";
+    private final String tenantId = "testTenantId";
+    private final MatchInfo matchInfo = MatchInfo.newBuilder()
+            .setTopicFilter("testTopicFilter")
+            .setReceiverId("testReceiverId")
+            .build();
+    private final DeliveryRequest request = DeliveryRequest.newBuilder()
+            .putPackage(tenantId, DeliveryPackage.newBuilder()
+                    .addPack(DeliveryPack.newBuilder()
+                            .addMatchInfo(matchInfo)
+                            .addMatchInfo(matchInfo).build())
+                    .build())
+            .build();
+
+    @BeforeMethod
+    public void setup() {
+        closeable = MockitoAnnotations.openMocks(this);
+        mqttBrokerClient = new MqttBrokerClient(rpcClient);
+    }
+
+    @AfterMethod
+    public void teardown() throws Exception {
+        closeable.close();
+    }
+
+    @Test
+    public void  deliveryPipelineWithOK() {
+        CompletableFuture<WriteReply> future = new CompletableFuture<>();
+        future.complete(WriteReply.newBuilder()
+                        .setReply(DeliveryReply.newBuilder()
+                                .putResult(tenantId, DeliveryResults.newBuilder()
+                                        .addResult(DeliveryResult.newBuilder()
+                                                .setCode(DeliveryResult.Code.OK)
+                                                .setMatchInfo(matchInfo)
+                                                .build())
+                                        .build())
+                                .build())
+                .build());
+
+        when(rpcClient.createRequestPipeline("", DeliveryGroupKeyUtil.parseServerId(delivererKey), "",
+                emptyMap(), OnlineInboxBrokerGrpc.getWriteMethod())).thenReturn(pipeline);
+        when(pipeline.invoke(any())).thenReturn(future);
+
+        deliverer = mqttBrokerClient.open(delivererKey);
+        DeliveryReply reply = deliverer.deliver(request).join();
+        DeliveryResults results = reply.getResultMap().get(tenantId);
+        assert  results != null;
+        DeliveryResult result = results.getResult(0);
+        assert result.getCode() == DeliveryResult.Code.OK;
+        assert result.getMatchInfo().equals(matchInfo);
+    }
+
+    @Test
+    public void deliveryPipelineWithSpecialException() {
+        CompletableFuture<WriteReply> future = new CompletableFuture<>();
+        future.completeExceptionally(new ServerNotFoundException("Test exception"));
+
+        when(rpcClient.createRequestPipeline("", DeliveryGroupKeyUtil.parseServerId(delivererKey), "",
+                emptyMap(), OnlineInboxBrokerGrpc.getWriteMethod())).thenReturn(pipeline);
+        when(pipeline.invoke(any())).thenReturn(future);
+
+        deliverer = mqttBrokerClient.open(delivererKey);
+        DeliveryReply reply = deliverer.deliver(request).join();
+        DeliveryResults results = reply.getResultMap().get(tenantId);
+        assert  results != null;
+        assert results.getResultList().size() == 1;
+        DeliveryResult result = results.getResult(0);
+        assert result.getCode() == DeliveryResult.Code.NO_SUB;
+        assert result.getMatchInfo().equals(matchInfo);
+    }
+}