From 0bb262dda3c7151431b1af9f8d02f1b77d0bf1e2 Mon Sep 17 00:00:00 2001
From: Sergey Matvienko <smatvienko@thingsboard.io>
Date: Wed, 22 Jan 2025 10:18:18 +0100
Subject: [PATCH] AbstractGatewaySessionHandler - add deregister callback for
 device creation futures on Gateway disconnect event

---
 .../AbstractGatewaySessionHandler.java        | 24 ++++++++++++++++++-
 1 file changed, 23 insertions(+), 1 deletion(-)

diff --git a/common/transport/mqtt/src/main/java/org/thingsboard/server/transport/mqtt/session/AbstractGatewaySessionHandler.java b/common/transport/mqtt/src/main/java/org/thingsboard/server/transport/mqtt/session/AbstractGatewaySessionHandler.java
index c2284ea76a..90ce8fccc2 100644
--- a/common/transport/mqtt/src/main/java/org/thingsboard/server/transport/mqtt/session/AbstractGatewaySessionHandler.java
+++ b/common/transport/mqtt/src/main/java/org/thingsboard/server/transport/mqtt/session/AbstractGatewaySessionHandler.java
@@ -16,8 +16,10 @@
 package org.thingsboard.server.transport.mqtt.session;
 
 import com.fasterxml.jackson.databind.JsonNode;
+import com.google.common.util.concurrent.FutureCallback;
 import com.google.common.util.concurrent.Futures;
 import com.google.common.util.concurrent.ListenableFuture;
+import com.google.common.util.concurrent.MoreExecutors;
 import com.google.common.util.concurrent.SettableFuture;
 import com.google.gson.JsonArray;
 import com.google.gson.JsonElement;
@@ -196,7 +198,27 @@ public abstract class AbstractGatewaySessionHandler<T extends AbstractGatewayDev
     }
 
     public void onDevicesDisconnect() {
-        devices.forEach(this::deregisterSession);
+        log.info("[{}] Gateway disconnect [{}]", gateway.getTenantId(), gateway.getDeviceId());
+        try {
+            deviceFutures.forEach((name, future) -> {
+                Futures.addCallback(future, new FutureCallback<T>() {
+                    @Override
+                    public void onSuccess(T result) {
+                        log.info("[{}] Gateway disconnect [{}] device deregister callback [{}]", gateway.getTenantId(), gateway.getDeviceId(), name);
+                        deregisterSession(name, result);
+                    }
+
+                    @Override
+                    public void onFailure(Throwable t) {
+
+                    }
+                }, MoreExecutors.directExecutor());
+            });
+
+            devices.forEach(this::deregisterSession);
+        } catch (Exception e) {
+            log.error("Gateway disconnect failure", e);
+        }
     }
 
     public void onDeviceDeleted(String deviceName) {