DefaultTbLocalSubscriptionService WS single lock refactored by tenantId; TbEntityLocalSubsInfo performance optimizations for remove subscription

application/src/main/java/org/thingsboard/server/service/subscription/DefaultTbLocalSubscriptionService.java

@@ -22,6 +22,7 @@ import org.springframework.beans.factory.annotation.Value;
 import org.springframework.context.annotation.Lazy;
 import org.springframework.context.event.EventListener;
 import org.springframework.stereotype.Service;
+import org.springframework.util.ConcurrentReferenceHashMap;
 import org.thingsboard.common.util.DeduplicationUtil;
 import org.thingsboard.common.util.DonAsynchron;
 import org.thingsboard.common.util.ThingsBoardExecutors;
@@ -120,12 +121,12 @@ public class DefaultTbLocalSubscriptionService implements TbLocalSubscriptionSer
     private String serviceId;
     private ExecutorService subscriptionUpdateExecutor;
 
-    private final Lock subsLock = new ReentrantLock();
+    private final ConcurrentReferenceHashMap<TenantId, Lock> locks = new ConcurrentReferenceHashMap<>(16, ConcurrentReferenceHashMap.ReferenceType.SOFT);
 
     @PostConstruct
     public void initExecutor() {
         subscriptionUpdateExecutor = ThingsBoardExecutors.newWorkStealingPool(20, getClass());
-        tsCallBackExecutor = Executors.newSingleThreadExecutor(ThingsBoardThreadFactory.forName("ts-sub-callback"));
+        tsCallBackExecutor = Executors.newFixedThreadPool(8, ThingsBoardThreadFactory.forName("ts-sub-callback")); //since we are using locks by TenantId
         serviceId = serviceInfoProvider.getServiceId();
         staleSessionCleanupExecutor = Executors.newSingleThreadScheduledExecutor(ThingsBoardThreadFactory.forName("stale-session-cleanup"));
         staleSessionCleanupExecutor.scheduleWithFixedDelay(this::cleanupStaleSessions, 60, 60, TimeUnit.SECONDS);
@@ -176,6 +177,10 @@ public class DefaultTbLocalSubscriptionService implements TbLocalSubscriptionSer
         });
     }
 
+    Lock getSubsLock(TenantId tenantId) {
+        return locks.computeIfAbsent(tenantId, x -> new ReentrantLock());
+    }
+
     @Override
     public void addSubscription(TbSubscription<?> subscription, WebSocketSessionRef sessionRef) {
         TenantId tenantId = subscription.getTenantId();
@@ -191,6 +196,7 @@ public class DefaultTbLocalSubscriptionService implements TbLocalSubscriptionSer
 
         log.debug("[{}][{}] Register subscription: {}", tenantId, entityId, subscription);
         SubscriptionModificationResult result;
+        final Lock subsLock = getSubsLock(tenantId);
         subsLock.lock();
         try {
             Map<Integer, TbSubscription<?>> sessionSubscriptions = subscriptionsBySessionId.computeIfAbsent(subscription.getSessionId(), k -> new ConcurrentHashMap<>());
@@ -219,14 +225,18 @@ public class DefaultTbLocalSubscriptionService implements TbLocalSubscriptionSer
         log.debug("[{}][{}] Processing sub event callback: {}.", entityId, seqNumber, entityUpdatesInfo);
         entityUpdates.put(entityId, entityUpdatesInfo);
         Set<TbSubscription<?>> pendingSubs = null;
-        subsLock.lock();
+        Lock subsLock = null;
         try {
             TbEntityLocalSubsInfo entitySubs = subscriptionsByEntityId.get(entityId);
             if (entitySubs != null) {
+                subsLock = getSubsLock(entitySubs.getTenantId());
+                subsLock.lock();
                 pendingSubs = entitySubs.clearPendingSubscriptions(seqNumber);
             }
         } finally {
-            subsLock.unlock();
+            if (subsLock != null) {
+                subsLock.unlock();
+            }
         }
         if (pendingSubs != null) {
             pendingSubs.forEach(this::checkMissedUpdates);
@@ -238,12 +248,14 @@ public class DefaultTbLocalSubscriptionService implements TbLocalSubscriptionSer
     public void cancelSubscription(String sessionId, int subscriptionId) {
         log.debug("[{}][{}] Going to remove subscription.", sessionId, subscriptionId);
         SubscriptionModificationResult result = null;
-        subsLock.lock();
+        Lock subsLock = null;
         try {
             Map<Integer, TbSubscription<?>> sessionSubscriptions = subscriptionsBySessionId.get(sessionId);
             if (sessionSubscriptions != null) {
                 TbSubscription<?> subscription = sessionSubscriptions.remove(subscriptionId);
                 if (subscription != null) {
+                    subsLock = getSubsLock(subscription.getTenantId());
+                    subsLock.lock();
                     if (sessionSubscriptions.isEmpty()) {
                         subscriptionsBySessionId.remove(sessionId);
                     }
@@ -255,7 +267,9 @@ public class DefaultTbLocalSubscriptionService implements TbLocalSubscriptionSer
                 log.debug("[{}] No session subscriptions found!", sessionId);
             }
         } finally {
-            subsLock.unlock();
+            if (subsLock != null) {
+                subsLock.unlock();
+            }
         }
         if (result != null && result.hasEvent()) {
             pushSubscriptionEvent(result);
@@ -266,18 +280,24 @@ public class DefaultTbLocalSubscriptionService implements TbLocalSubscriptionSer
     public void cancelAllSessionSubscriptions(String sessionId) {
         log.debug("[{}] Going to remove session subscriptions.", sessionId);
         List<SubscriptionModificationResult> results = new ArrayList<>();
-        subsLock.lock();
+        Lock subsLock = null;
         try {
             Map<Integer, TbSubscription<?>> sessionSubscriptions = subscriptionsBySessionId.remove(sessionId);
             if (sessionSubscriptions != null) {
                 for (TbSubscription<?> subscription : sessionSubscriptions.values()) {
+                    if (subsLock == null) {
+                    subsLock = getSubsLock(subscription.getTenantId());
+                    subsLock.lock();
+                    }
                     results.add(modifySubscription(subscription.getTenantId(), subscription.getEntityId(), subscription, false));
                 }
             } else {
                 log.debug("[{}] No session subscriptions found!", sessionId);
             }
         } finally {
-            subsLock.unlock();
+            if (subsLock != null) {
+                subsLock.unlock();
+            }
         }
         results.stream().filter(SubscriptionModificationResult::hasEvent).forEach(this::pushSubscriptionEvent);
     }

application/src/main/java/org/thingsboard/server/service/subscription/TbAlarmsSubscription.java

@@ -36,12 +36,8 @@ public class TbAlarmsSubscription extends TbSubscription<AlarmSubscriptionUpdate
     }
 
     @Override
-    public boolean equals(Object o) {
-        return super.equals(o);
+    protected boolean canEqual(final Object other) {
+        return other instanceof TbAlarmsSubscription;
     }
 
-    @Override
-    public int hashCode() {
-        return super.hashCode();
-    }
 }

application/src/main/java/org/thingsboard/server/service/subscription/TbAttributeSubscription.java

@@ -43,12 +43,8 @@ public class TbAttributeSubscription extends TbSubscription<TelemetrySubscriptio
     }
 
     @Override
-    public boolean equals(Object o) {
-        return super.equals(o);
+    protected boolean canEqual(final Object other) {
+        return other instanceof TbAttributeSubscription;
     }
 
-    @Override
-    public int hashCode() {
-        return super.hashCode();
-    }
 }

application/src/main/java/org/thingsboard/server/service/subscription/TbEntityLocalSubsInfo.java

@@ -27,8 +27,6 @@ import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.locks.Lock;
-import java.util.concurrent.locks.ReentrantLock;
 
 /**
  * Information about the local websocket subscriptions.
@@ -42,8 +40,6 @@ public class TbEntityLocalSubsInfo {
     @Getter
     private final EntityId entityId;
     @Getter
-    private final Lock lock = new ReentrantLock();
-    @Getter
     private final Set<TbSubscription<?>> subs = ConcurrentHashMap.newKeySet();
     private volatile TbSubscriptionsInfo state = new TbSubscriptionsInfo();
 
@@ -88,6 +84,7 @@ public class TbEntityLocalSubsInfo {
                 if (!newState.attrAllKeys) {
                     if (attrSub.isAllKeys()) {
                         newState.attrAllKeys = true;
+                        newState.attrKeys = null;
                         stateChanged = true;
                     } else {
                         if (newState.attrKeys == null) {
@@ -104,6 +101,7 @@ public class TbEntityLocalSubsInfo {
                 if (!newState.tsAllKeys) {
                     if (tsSub.isAllKeys()) {
                         newState.tsAllKeys = true;
+                        newState.tsKeys = null;
                         stateChanged = true;
                     } else {
                         if (newState.tsKeys == null) {
@@ -137,8 +135,30 @@ public class TbEntityLocalSubsInfo {
             return toEvent(ComponentLifecycleEvent.DELETED);
         }
         TbSubscriptionsInfo oldState = state.copy();
-        TbSubscriptionsInfo newState = new TbSubscriptionsInfo();
+
+        //copy unchanged state only
+        TbSubscriptionsInfo newState = state.copy();
+        switch (sub.getType()) {
+            case NOTIFICATIONS:
+            case NOTIFICATIONS_COUNT:
+                newState.notifications = false;
+                break;
+            case ALARMS:
+                newState.alarms = false;
+                break;
+            case ATTRIBUTES:
+                newState.attrAllKeys = false;
+                newState.attrKeys = null;
+                break;
+            case TIMESERIES:
+                newState.tsAllKeys = false;
+                newState.tsKeys = null;
+        }
+
         for (TbSubscription<?> subscription : subs) {
+            if (subscription.getType() != sub.getType()) {
+                continue; // skip unchanged types
+            }
             switch (subscription.getType()) {
                 case NOTIFICATIONS:
                 case NOTIFICATIONS_COUNT:
@@ -155,6 +175,7 @@ public class TbEntityLocalSubsInfo {
                     var attrSub = (TbAttributeSubscription) subscription;
                     if (!newState.attrAllKeys && attrSub.isAllKeys()) {
                         newState.attrAllKeys = true;
+                        newState.attrKeys = null;
                         continue;
                     }
                     if (newState.attrKeys == null) {
@@ -167,6 +188,7 @@ public class TbEntityLocalSubsInfo {
                     var tsSub = (TbTimeSeriesSubscription) subscription;
                     if (!newState.tsAllKeys && tsSub.isAllKeys()) {
                         newState.tsAllKeys = true;
+                        newState.tsKeys = null;
                         continue;
                     }
                     if (newState.tsKeys == null) {

application/src/main/java/org/thingsboard/server/service/subscription/TbSubscription.java

@@ -15,41 +15,25 @@
  */
 package org.thingsboard.server.service.subscription;
 
-import lombok.AllArgsConstructor;
 import lombok.Data;
+import lombok.EqualsAndHashCode;
 import org.thingsboard.server.common.data.id.EntityId;
 import org.thingsboard.server.common.data.id.TenantId;
 
-import java.util.Objects;
-import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.BiConsumer;
 
 @Data
-@AllArgsConstructor
+@EqualsAndHashCode(cacheStrategy = EqualsAndHashCode.CacheStrategy.LAZY, doNotUseGetters = true)
 public abstract class TbSubscription<T> {
 
+    @EqualsAndHashCode.Exclude
     private final String serviceId;
     private final String sessionId;
     private final int subscriptionId;
     private final TenantId tenantId;
     private final EntityId entityId;
     private final TbSubscriptionType type;
+    @EqualsAndHashCode.Exclude
     private final BiConsumer<TbSubscription<T>, T> updateProcessor;
 
-    @Override
-    public boolean equals(Object o) {
-        if (this == o) return true;
-        if (o == null || getClass() != o.getClass()) return false;
-        TbSubscription that = (TbSubscription) o;
-        return subscriptionId == that.subscriptionId &&
-                sessionId.equals(that.sessionId) &&
-                tenantId.equals(that.tenantId) &&
-                entityId.equals(that.entityId) &&
-                type == that.type;
-    }
-
-    @Override
-    public int hashCode() {
-        return Objects.hash(sessionId, subscriptionId, tenantId, entityId, type);
-    }
 }

application/src/main/java/org/thingsboard/server/service/subscription/TbSubscriptionUtils.java

@@ -107,8 +107,9 @@ public class TbSubscriptionUtils {
                 .type(event);
         if (!ComponentLifecycleEvent.DELETED.equals(event)) {
             builder.info(new TbSubscriptionsInfo(proto.getNotifications(), proto.getAlarms(),
-                    proto.getTsAllKeys(), proto.getTsKeysCount() > 0 ? new HashSet<>(proto.getTsKeysList()) : null,
-                    proto.getAttrAllKeys(), proto.getAttrKeysCount() > 0 ? new HashSet<>(proto.getAttrKeysList()) : null,
+                    proto.getTsAllKeys(), proto.getAttrAllKeys(),
+                    proto.getTsKeysCount() > 0 ? new HashSet<>(proto.getTsKeysList()) : null,
+                    proto.getAttrKeysCount() > 0 ? new HashSet<>(proto.getAttrKeysList()) : null,
                     proto.getSeqNumber()));
         }
         return builder.build();

application/src/main/java/org/thingsboard/server/service/subscription/TbSubscriptionsInfo.java

@@ -20,25 +20,24 @@ import lombok.EqualsAndHashCode;
 import lombok.RequiredArgsConstructor;
 import lombok.ToString;
 
-import java.util.HashSet;
 import java.util.Set;
-import java.util.concurrent.ConcurrentHashMap;
 
 /**
  * Information about the local websocket subscriptions.
  */
 @RequiredArgsConstructor
 @AllArgsConstructor
-@EqualsAndHashCode(exclude = {"seqNumber"})
+@EqualsAndHashCode
 @ToString
 public class TbSubscriptionsInfo {
 
     protected boolean notifications;
     protected boolean alarms;
     protected boolean tsAllKeys;
+    protected boolean attrAllKeys; // primitives first for equals performance
     protected Set<String> tsKeys;
-    protected boolean attrAllKeys;
     protected Set<String> attrKeys;
+    @EqualsAndHashCode.Exclude
     protected int seqNumber;
 
     public boolean isEmpty() {
@@ -50,7 +49,7 @@ public class TbSubscriptionsInfo {
     }
 
     protected TbSubscriptionsInfo copy(int seqNumber) {
-        return new TbSubscriptionsInfo(notifications, alarms, tsAllKeys, tsKeys, attrAllKeys, attrKeys, seqNumber);
+        return new TbSubscriptionsInfo(notifications, alarms, tsAllKeys, attrAllKeys, tsKeys, attrKeys, seqNumber);
     }
 
 }

application/src/main/java/org/thingsboard/server/service/subscription/TbTimeSeriesSubscription.java

@@ -53,12 +53,8 @@ public class TbTimeSeriesSubscription extends TbSubscription<TelemetrySubscripti
     }
 
     @Override
-    public boolean equals(Object o) {
-        return super.equals(o);
+    protected boolean canEqual(final Object other) {
+        return other instanceof TbTimeSeriesSubscription;
     }
 
-    @Override
-    public int hashCode() {
-        return super.hashCode();
-    }
 }

application/src/main/java/org/thingsboard/server/service/ws/notification/sub/AbstractNotificationSubscription.java

@@ -35,4 +35,9 @@ public abstract class AbstractNotificationSubscription<T> extends TbSubscription
         super(serviceId, sessionId, subscriptionId, tenantId, entityId, type, updateProcessor);
     }
 
+    @Override
+    protected boolean canEqual(final Object other) {
+        return other instanceof AbstractNotificationSubscription;
+    }
+
 }

application/src/main/java/org/thingsboard/server/service/ws/notification/sub/NotificationsCountSubscription.java

@@ -23,7 +23,6 @@ import org.thingsboard.server.service.subscription.TbSubscription;
 import org.thingsboard.server.service.subscription.TbSubscriptionType;
 import org.thingsboard.server.service.ws.notification.cmd.UnreadNotificationsCountUpdate;
 
-import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.BiConsumer;
 
 @Getter
@@ -35,6 +34,11 @@ public class NotificationsCountSubscription extends AbstractNotificationSubscrip
         super(serviceId, sessionId, subscriptionId, tenantId, entityId, TbSubscriptionType.NOTIFICATIONS_COUNT, updateProcessor);
     }
 
+    @Override
+    protected boolean canEqual(final Object other) {
+        return other instanceof NotificationsCountSubscription;
+    }
+
     public UnreadNotificationsCountUpdate createUpdate() {
         return UnreadNotificationsCountUpdate.builder()
                 .cmdId(getSubscriptionId())

application/src/main/java/org/thingsboard/server/service/ws/notification/sub/NotificationsSubscription.java

@@ -30,7 +30,6 @@ import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.UUID;
-import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.BiConsumer;
 import java.util.stream.Collectors;
 
@@ -48,6 +47,11 @@ public class NotificationsSubscription extends AbstractNotificationSubscription<
         this.limit = limit;
     }
 
+    @Override
+    protected boolean canEqual(final Object other) {
+        return other instanceof NotificationsSubscription;
+    }
+
     public UnreadNotificationsUpdate createFullUpdate() {
         return UnreadNotificationsUpdate.builder()
                 .cmdId(getSubscriptionId())