Refactoring to avoid race conditions on device registration

common/transport/lwm2m/src/main/java/org/thingsboard/server/transport/lwm2m/server/DefaultLwM2MTransportMsgHandler.java

@@ -58,6 +58,8 @@ import org.thingsboard.server.queue.util.TbLwM2mTransportComponent;
 import org.thingsboard.server.transport.lwm2m.config.LwM2MTransportServerConfig;
 import org.thingsboard.server.transport.lwm2m.server.LwM2mTransportUtil.LwM2mTypeOper;
 import org.thingsboard.server.transport.lwm2m.server.adaptors.LwM2MJsonAdaptor;
+import org.thingsboard.server.transport.lwm2m.server.client.LwM2MClientState;
+import org.thingsboard.server.transport.lwm2m.server.client.LwM2MClientStateException;
 import org.thingsboard.server.transport.lwm2m.server.client.LwM2mClient;
 import org.thingsboard.server.transport.lwm2m.server.client.LwM2mClientContext;
 import org.thingsboard.server.transport.lwm2m.server.client.LwM2mClientProfile;
@@ -91,7 +93,6 @@ import static org.thingsboard.server.common.data.ota.OtaPackageUpdateStatus.DOWN
 import static org.thingsboard.server.common.data.ota.OtaPackageUpdateStatus.UPDATING;
 import static org.thingsboard.server.common.data.lwm2m.LwM2mConstants.LWM2M_SEPARATOR_PATH;
 import static org.thingsboard.server.transport.lwm2m.server.LwM2mTransportServerHelper.getValueFromKvProto;
-import static org.thingsboard.server.transport.lwm2m.server.LwM2mTransportUtil.CLIENT_NOT_AUTHORIZED;
 import static org.thingsboard.server.transport.lwm2m.server.LwM2mTransportUtil.DEVICE_ATTRIBUTES_REQUEST;
 import static org.thingsboard.server.transport.lwm2m.server.LwM2mTransportUtil.FW_ID;
 import static org.thingsboard.server.transport.lwm2m.server.LwM2mTransportUtil.FW_RESULT_ID;
@@ -184,9 +185,10 @@ public class DefaultLwM2MTransportMsgHandler implements LwM2mTransportMsgHandler
         registrationExecutor.submit(() -> {
             try {
                 log.warn("[{}] [{{}] Client: create after Registration", registration.getEndpoint(), registration.getId());
-                LwM2mClient lwM2MClient = this.clientContext.registerOrUpdate(registration);
+                LwM2mClient lwM2MClient = this.clientContext.getClientByEndpoint(registration.getEndpoint());
                 if (lwM2MClient != null) {
-                    SessionInfoProto sessionInfo = this.getSessionInfoOrCloseSession(lwM2MClient);
+                    this.clientContext.register(lwM2MClient, registration);
+                    SessionInfoProto sessionInfo = lwM2MClient.getSession();
                     if (sessionInfo != null) {
                         transportService.registerAsyncSession(sessionInfo, new LwM2mSessionMsgListener(this, sessionInfo));
                         TransportProtos.TransportToDeviceActorMsg msg = TransportProtos.TransportToDeviceActorMsg.newBuilder()
@@ -199,13 +201,19 @@ public class DefaultLwM2MTransportMsgHandler implements LwM2mTransportMsgHandler
                         this.getInfoFirmwareUpdate(lwM2MClient, null);
                         this.getInfoSoftwareUpdate(lwM2MClient, null);
                         this.initLwM2mFromClientValue(registration, lwM2MClient);
-                        this.sendLogsToThingsboard(LOG_LW2M_INFO + ": Client create after Registration", registration.getId());
+                        this.sendLogsToThingsboard(lwM2MClient, LOG_LW2M_INFO + ": Client create after Registration");
                     } else {
                         log.error("Client: [{}] onRegistered [{}] name  [{}] sessionInfo ", registration.getId(), registration.getEndpoint(), null);
                     }
                 } else {
                     log.error("Client: [{}] onRegistered [{}] name  [{}] lwM2MClient ", registration.getId(), registration.getEndpoint(), null);
                 }
+            } catch (LwM2MClientStateException stateException) {
+                if (LwM2MClientState.UNREGISTERED.equals(stateException.getState())) {
+                    log.info("[{}] retry registration due to race condition: [{}].", registration.getEndpoint(), stateException.getState());
+                    // Race condition detected and the client was in progress of unregistration while new registration arrived. Let's try again.
+                    onRegistered(registration, previousObservations);
+                }
             } catch (Throwable t) {
                 log.error("[{}] endpoint [{}] error Unable registration.", registration.getEndpoint(), t);
             }
@@ -219,25 +227,26 @@ public class DefaultLwM2MTransportMsgHandler implements LwM2mTransportMsgHandler
      */
     public void updatedReg(Registration registration) {
         updateRegistrationExecutor.submit(() -> {
+            LwM2mClient lwM2MClient = clientContext.getClientByEndpoint(registration.getEndpoint());
             try {
-                LwM2mClient client = clientContext.getOrRegister(registration);
+                clientContext.updateRegistration(lwM2MClient, registration);
-                if (client != null && client.getSession() != null) {
+                TransportProtos.SessionInfoProto sessionInfo = lwM2MClient.getSession();
-                    SessionInfoProto sessionInfo = client.getSession();
+                this.reportActivityAndRegister(sessionInfo);
-                    this.reportActivityAndRegister(sessionInfo);
+                if (registration.getQueueMode()) {
-                    if (registration.getQueueMode()) {
+                    LwM2mQueuedRequest request;
-                        LwM2mQueuedRequest request;
+                    while ((request = lwM2MClient.getQueuedRequests().poll()) != null) {
-                        while ((request = client.getQueuedRequests().poll()) != null) {
+                        request.send();
-                            request.send();
-                        }
                     }
-                    this.sendLogsToThingsboard(LOG_LW2M_INFO + ": Client update Registration", registration.getId());
+                }
+            } catch (LwM2MClientStateException stateException) {
+                if (LwM2MClientState.UNREGISTERED.equals(stateException.getState())) {
+                    log.info("[{}] update registration failed because client was already unregistered: [{}].", registration.getEndpoint(), stateException.getState());
                 } else {
-                    log.error("Client: [{}] updatedReg [{}] name  [{}] sessionInfo ", registration.getId(), registration.getEndpoint(), null);
+                    log.info("[{}] update registration: [{}] {}.", registration.getEndpoint(), stateException.getState(), stateException.getMessage());
-                    this.sendLogsToThingsboard(LOG_LW2M_ERROR + ": Client update Registration", registration.getId());
                 }
             } catch (Throwable t) {
                 log.error("[{}] endpoint [{}] error Unable update registration.", registration.getEndpoint(), t);
-                this.sendLogsToThingsboard(LOG_LW2M_ERROR + String.format(": Client update Registration, %s", t.getMessage()), registration.getId());
+                this.sendLogsToThingsboard(lwM2MClient, LOG_LW2M_ERROR + String.format(": Client update Registration, %s", t.getMessage()));
             }
         });
     }
@@ -248,34 +257,32 @@ public class DefaultLwM2MTransportMsgHandler implements LwM2mTransportMsgHandler
      */
     public void unReg(Registration registration, Collection<Observation> observations) {
         unRegistrationExecutor.submit(() -> {
+            LwM2mClient client = clientContext.getClientByEndpoint(registration.getEndpoint());
             try {
-                this.sendLogsToThingsboard(LOG_LW2M_INFO + ": Client unRegistration", registration.getId());
+                this.sendLogsToThingsboard(client, LOG_LW2M_INFO + ": Client unRegistration");
-                this.closeClientSession(registration);
+                clientContext.unregister(client, registration);
+                SessionInfoProto sessionInfo = client.getSession();
+                if (sessionInfo != null) {
+                    transportService.deregisterSession(sessionInfo);
+                    sessionStore.remove(registration.getEndpoint());
+                    this.doCloseSession(sessionInfo);
+                    log.info("Client close session: [{}] unReg [{}] name  [{}] profile ", registration.getId(), registration.getEndpoint(), sessionInfo.getDeviceType());
+                } else {
+                    log.error("Client close session: [{}] unReg [{}] name  [{}] sessionInfo ", registration.getId(), registration.getEndpoint(), null);
+                }
+            } catch (LwM2MClientStateException stateException) {
+                log.info("[{}] delete registration: [{}] {}.", registration.getEndpoint(), stateException.getState(), stateException.getMessage());
             } catch (Throwable t) {
                 log.error("[{}] endpoint [{}] error Unable un registration.", registration.getEndpoint(), t);
-                this.sendLogsToThingsboard(LOG_LW2M_ERROR + String.format(": Client Unable un Registration, %s", t.getMessage()), registration.getId());
+                this.sendLogsToThingsboard(client, LOG_LW2M_ERROR + String.format(": Client Unable un Registration, %s", t.getMessage()));
             }
         });
     }
 
-    private void closeClientSession(Registration registration) {
-        SessionInfoProto sessionInfo = this.getSessionInfoOrCloseSession(registration);
-        if (sessionInfo != null) {
-            transportService.deregisterSession(sessionInfo);
-            sessionStore.remove(registration.getEndpoint());
-            this.doCloseSession(sessionInfo);
-            clientContext.removeClientByRegistrationId(registration.getId());
-            log.info("Client close session: [{}] unReg [{}] name  [{}] profile ", registration.getId(), registration.getEndpoint(), sessionInfo.getDeviceType());
-        } else {
-            log.error("Client close session: [{}] unReg [{}] name  [{}] sessionInfo ", registration.getId(), registration.getEndpoint(), null);
-        }
-    }
-
     @Override
     public void onSleepingDev(Registration registration) {
         log.info("[{}] [{}] Received endpoint Sleeping version event", registration.getId(), registration.getEndpoint());
-        this.sendLogsToThingsboard(LOG_LW2M_INFO + ": Client is sleeping!", registration.getId());
+        this.sendLogsToThingsboard(clientContext.getClientByEndpoint(registration.getEndpoint()), LOG_LW2M_INFO + ": Client is sleeping!");
-
         //TODO: associate endpointId with device information.
     }
 
@@ -300,7 +307,7 @@ public class DefaultLwM2MTransportMsgHandler implements LwM2mTransportMsgHandler
     @Override
     public void onUpdateValueAfterReadResponse(Registration registration, String path, ReadResponse response, Lwm2mClientRpcRequest rpcRequest) {
         if (response.getContent() != null) {
-            LwM2mClient lwM2MClient = clientContext.getOrRegister(registration);
+            LwM2mClient lwM2MClient = clientContext.getClientByEndpoint(registration.getEndpoint());
             ObjectModel objectModelVersion = lwM2MClient.getObjectModel(path, this.config.getModelProvider());
             if (objectModelVersion != null) {
                 if (response.getContent() instanceof LwM2mObject) {
@@ -332,7 +339,7 @@ public class DefaultLwM2MTransportMsgHandler implements LwM2mTransportMsgHandler
         }
         String msg = String.format("%s: type operation %s path - %s value - %s", LOG_LW2M_INFO,
                 READ, pathIdVer, value);
-        this.sendLogsToThingsboard(msg, registration.getId());
+        this.sendLogsToThingsboard(lwM2MClient, msg);
         rpcRequest.setValueMsg(String.format("%s", value));
         this.sentRpcResponse(rpcRequest, response.getCode().getName(), (String) value, LOG_LW2M_VALUE);
     }
@@ -352,7 +359,7 @@ public class DefaultLwM2MTransportMsgHandler implements LwM2mTransportMsgHandler
     public void onAttributeUpdate(AttributeUpdateNotificationMsg msg, TransportProtos.SessionInfoProto sessionInfo) {
         LwM2mClient lwM2MClient = clientContext.getClient(sessionInfo);
         if (msg.getSharedUpdatedCount() > 0 && lwM2MClient != null) {
-            log.warn ("2) OnAttributeUpdate, SharedUpdatedList() [{}]", msg.getSharedUpdatedList());
+            log.warn("2) OnAttributeUpdate, SharedUpdatedList() [{}]", msg.getSharedUpdatedList());
             msg.getSharedUpdatedList().forEach(tsKvProto -> {
                 String pathName = tsKvProto.getKv().getKey();
                 String pathIdVer = this.getPresentPathIntoProfile(sessionInfo, pathName);
@@ -377,13 +384,13 @@ public class DefaultLwM2MTransportMsgHandler implements LwM2mTransportMsgHandler
                         log.error("Resource path - [{}] value - [{}] is not Writable and cannot be updated", pathIdVer, valueNew);
                         String logMsg = String.format("%s: attributeUpdate: Resource path - %s value - %s is not Writable and cannot be updated",
                                 LOG_LW2M_ERROR, pathIdVer, valueNew);
-                        this.sendLogsToThingsboard(logMsg, lwM2MClient.getRegistration().getId());
+                        this.sendLogsToThingsboard(lwM2MClient, logMsg);
                     }
                 } else if (!isFwSwWords(pathName)) {
                     log.error("Resource name name - [{}] value - [{}] is not present as attribute/telemetry in profile and cannot be updated", pathName, valueNew);
                     String logMsg = String.format("%s: attributeUpdate: attribute name - %s value - %s is not present as attribute in profile and cannot be updated",
                             LOG_LW2M_ERROR, pathName, valueNew);
-                    this.sendLogsToThingsboard(logMsg, lwM2MClient.getRegistration().getId());
+                    this.sendLogsToThingsboard(lwM2MClient, logMsg);
                 }
 
             });
@@ -396,9 +403,8 @@ public class DefaultLwM2MTransportMsgHandler implements LwM2mTransportMsgHandler
                 }
             });
             log.info("[{}] delete [{}]  onAttributeUpdate", msg.getSharedDeletedList(), sessionInfo);
-        }
+        } else if (lwM2MClient == null) {
-        else if (lwM2MClient == null) {
+            log.error("OnAttributeUpdate, lwM2MClient is null");
-            log.error ("OnAttributeUpdate, lwM2MClient is null");
         }
     }
 
@@ -408,12 +414,11 @@ public class DefaultLwM2MTransportMsgHandler implements LwM2mTransportMsgHandler
      */
     @Override
     public void onDeviceProfileUpdate(SessionInfoProto sessionInfo, DeviceProfile deviceProfile) {
-        Set<LwM2mClient> clients = clientContext.getLwM2mClients()
+        List<LwM2mClient> clients = clientContext.getLwM2mClients()
-                .stream().filter(e -> e.getProfileId().equals(deviceProfile.getUuidId())).collect(Collectors.toSet());
+                .stream().filter(e -> e.getProfileId().equals(deviceProfile.getUuidId())).collect(Collectors.toList());
         clients.forEach(client -> client.onDeviceProfileUpdate(deviceProfile));
-        Set<String> registrationIds = clients.stream().map(LwM2mClient::getRegistration).map(Registration::getId).collect(Collectors.toSet());
+        if (clients.size() > 0) {
-        if (registrationIds.size() > 0) {
+            this.onDeviceProfileUpdate(clients, deviceProfile);
-            this.onDeviceProfileUpdate(registrationIds, deviceProfile);
         }
     }
 
@@ -446,7 +451,7 @@ public class DefaultLwM2MTransportMsgHandler implements LwM2mTransportMsgHandler
     public void onToDeviceRpcRequest(TransportProtos.ToDeviceRpcRequestMsg toDeviceRpcRequestMsg, SessionInfoProto sessionInfo) {
         // #1
         this.checkRpcRequestTimeout();
-        log.warn ("4) toDeviceRpcRequestMsg: [{}], sessionUUID: [{}]", toDeviceRpcRequestMsg,  new UUID(sessionInfo.getSessionIdMSB(), sessionInfo.getSessionIdLSB()));
+        log.warn("4) toDeviceRpcRequestMsg: [{}], sessionUUID: [{}]", toDeviceRpcRequestMsg, new UUID(sessionInfo.getSessionIdMSB(), sessionInfo.getSessionIdLSB()));
         String bodyParams = StringUtils.trimToNull(toDeviceRpcRequestMsg.getParams()) != null ? toDeviceRpcRequestMsg.getParams() : "null";
         LwM2mTypeOper lwM2mTypeOper = setValidTypeOper(toDeviceRpcRequestMsg.getMethodName());
         UUID requestUUID = new UUID(toDeviceRpcRequestMsg.getRequestIdMSB(), toDeviceRpcRequestMsg.getRequestIdLSB());
@@ -506,7 +511,7 @@ public class DefaultLwM2MTransportMsgHandler implements LwM2mTransportMsgHandler
 
     @Override
     public void onToDeviceRpcResponse(TransportProtos.ToDeviceRpcResponseMsg toDeviceResponse, SessionInfoProto sessionInfo) {
-        log.warn ("5) onToDeviceRpcResponse: [{}], sessionUUID: [{}]", toDeviceResponse,  new UUID(sessionInfo.getSessionIdMSB(), sessionInfo.getSessionIdLSB()));
+        log.warn("5) onToDeviceRpcResponse: [{}], sessionUUID: [{}]", toDeviceResponse, new UUID(sessionInfo.getSessionIdMSB(), sessionInfo.getSessionIdLSB()));
         transportService.process(sessionInfo, toDeviceResponse, null);
     }
 
@@ -558,7 +563,7 @@ public class DefaultLwM2MTransportMsgHandler implements LwM2mTransportMsgHandler
     @Override
     public void onAwakeDev(Registration registration) {
         log.trace("[{}] [{}] Received endpoint Awake version event", registration.getId(), registration.getEndpoint());
-        this.sendLogsToThingsboard(LOG_LW2M_INFO + ": Client is awake!", registration.getId());
+        this.sendLogsToThingsboard(clientContext.getClientByEndpoint(registration.getEndpoint()), LOG_LW2M_INFO + ": Client is awake!");
         //TODO: associate endpointId with device information.
     }
 
@@ -567,13 +572,17 @@ public class DefaultLwM2MTransportMsgHandler implements LwM2mTransportMsgHandler
      * @param registrationId - Id of Registration LwM2M Client
      */
     @Override
-    public void sendLogsToThingsboard(String logMsg, String registrationId) {
+    public void sendLogsToThingsboard2(String registrationId, String logMsg) {
-        SessionInfoProto sessionInfo = this.getSessionInfoOrCloseSession(registrationId);
+        sendLogsToThingsboard(clientContext.getClientByRegistrationId(registrationId), logMsg);
-        if (logMsg != null && sessionInfo != null) {
+    }
+
+    @Override
+    public void sendLogsToThingsboard(LwM2mClient client, String logMsg) {
+        if (logMsg != null && client != null && client.getSession() != null) {
             if (logMsg.length() > 1024) {
                 logMsg = logMsg.substring(0, 1024);
             }
-            this.helper.sendParametersOnThingsboardTelemetry(this.helper.getKvStringtoThingsboard(LOG_LW2M_TELEMETRY, logMsg), sessionInfo);
+            this.helper.sendParametersOnThingsboardTelemetry(this.helper.getKvStringtoThingsboard(LOG_LW2M_TELEMETRY, logMsg), client.getSession());
         }
     }
 
@@ -645,7 +654,7 @@ public class DefaultLwM2MTransportMsgHandler implements LwM2mTransportMsgHandler
      * @param path          - resource
      */
     private void updateResourcesValue(Registration registration, LwM2mResource lwM2mResource, String path) {
-        LwM2mClient lwM2MClient = clientContext.getOrRegister(registration);
+        LwM2mClient lwM2MClient = clientContext.getClientByEndpoint(registration.getEndpoint());
         if (lwM2MClient.saveResourceValue(path, lwM2mResource, this.config.getModelProvider())) {
             /** version != null
              * set setClient_fw_info... = value
@@ -798,7 +807,7 @@ public class DefaultLwM2MTransportMsgHandler implements LwM2mTransportMsgHandler
     }
 
     private void onDeviceUpdate(LwM2mClient lwM2MClient, Device device, Optional<DeviceProfile> deviceProfileOpt) {
-        deviceProfileOpt.ifPresent(deviceProfile -> this.onDeviceProfileUpdate(Collections.singleton(lwM2MClient.getRegistration().getId()), deviceProfile));
+        deviceProfileOpt.ifPresent(deviceProfile -> this.onDeviceProfileUpdate(Collections.singletonList(lwM2MClient), deviceProfile));
         lwM2MClient.onDeviceUpdate(device, deviceProfileOpt);
     }
 
@@ -843,7 +852,7 @@ public class DefaultLwM2MTransportMsgHandler implements LwM2mTransportMsgHandler
     }
 
     private TransportProtos.KeyValueProto getKvToThingsboard(String pathIdVer, Registration registration) {
-        LwM2mClient lwM2MClient = this.clientContext.getClientByRegistrationId(registration.getId());
+        LwM2mClient lwM2MClient = this.clientContext.getClientByEndpoint(registration.getEndpoint());
         JsonObject names = clientContext.getProfiles().get(lwM2MClient.getProfileId()).getPostKeyNameProfile();
         if (names != null && names.has(pathIdVer)) {
             String resourceName = names.get(pathIdVer).getAsString();
@@ -892,9 +901,7 @@ public class DefaultLwM2MTransportMsgHandler implements LwM2mTransportMsgHandler
             ResourceModel.Type expectedType = this.helper.getResourceModelTypeEqualsKvProtoValueType(currentType, pathIdVer);
             return this.converter.convertValue(resourceValue.getValue(), currentType, expectedType,
                     new LwM2mPath(convertPathFromIdVerToObjectId(pathIdVer)));
-        }
+        } else {
-
-        else {
             return null;
         }
     }
@@ -955,10 +962,10 @@ public class DefaultLwM2MTransportMsgHandler implements LwM2mTransportMsgHandler
      * #6.1 - update WriteAttribute
      * #6.2 - del WriteAttribute
      *
-     * @param registrationIds -
+     * @param clients       -
-     * @param deviceProfile   -
+     * @param deviceProfile -
      */
-    private void onDeviceProfileUpdate(Set<String> registrationIds, DeviceProfile deviceProfile) {
+    private void onDeviceProfileUpdate(List<LwM2mClient> clients, DeviceProfile deviceProfile) {
         LwM2mClientProfile lwM2MClientProfileOld = clientContext.getProfiles().get(deviceProfile.getUuidId()).clone();
         if (clientContext.profileUpdate(deviceProfile) != null) {
             // #1
@@ -1009,15 +1016,14 @@ public class DefaultLwM2MTransportMsgHandler implements LwM2mTransportMsgHandler
 
             // #3.4, #6
             if (!attributeLwm2mOld.equals(attributeLwm2mNew)) {
-                this.getAnalyzerAttributeLwm2m(registrationIds, attributeLwm2mOld, attributeLwm2mNew);
+                this.getAnalyzerAttributeLwm2m(clients, attributeLwm2mOld, attributeLwm2mNew);
             }
 
             // #4.1 add
             if (sendAttrToThingsboard.getPathPostParametersAdd().size() > 0) {
                 // update value in Resources
-                registrationIds.forEach(registrationId -> {
+                clients.forEach(client -> {
-                    Registration registration = clientContext.getRegistration(registrationId);
+                    this.readObserveFromProfile(client.getRegistration(), sendAttrToThingsboard.getPathPostParametersAdd(), READ);
-                    this.readObserveFromProfile(registration, sendAttrToThingsboard.getPathPostParametersAdd(), READ);
                 });
             }
             // #4.2 del
@@ -1041,8 +1047,8 @@ public class DefaultLwM2MTransportMsgHandler implements LwM2mTransportMsgHandler
                 // does not include oldObserve
                 ResultsAnalyzerParameters postObserveAnalyzer = this.getAnalyzerParameters(sendObserveToClientOld.getPathPostParametersAdd(), sendObserveToClientNew.getPathPostParametersAdd());
                 //  send Request observe to Client
-                registrationIds.forEach(registrationId -> {
+                clients.forEach(client -> {
-                    Registration registration = clientContext.getRegistration(registrationId);
+                    Registration registration = client.getRegistration();
                     if (postObserveAnalyzer.getPathPostParametersAdd().size() > 0) {
                         this.readObserveFromProfile(registration, postObserveAnalyzer.getPathPostParametersAdd(), OBSERVE);
                     }
@@ -1124,7 +1130,7 @@ public class DefaultLwM2MTransportMsgHandler implements LwM2mTransportMsgHandler
      * @param attributeLwm2mNew -
      * @return
      */
-    private void getAnalyzerAttributeLwm2m(Set<String> registrationIds, JsonObject attributeLwm2mOld, JsonObject attributeLwm2mNew) {
+    private void getAnalyzerAttributeLwm2m(List<LwM2mClient> clients, JsonObject attributeLwm2mOld, JsonObject attributeLwm2mNew) {
         ResultsAnalyzerParameters analyzerParameters = new ResultsAnalyzerParameters();
         ConcurrentHashMap<String, Object> lwm2mAttributesOld = new Gson().fromJson(attributeLwm2mOld.toString(),
                 new TypeToken<ConcurrentHashMap<String, Object>>() {
@@ -1146,8 +1152,8 @@ public class DefaultLwM2MTransportMsgHandler implements LwM2mTransportMsgHandler
         // #6
         // #6.2
         if (analyzerParameters.getPathPostParametersAdd().size() > 0) {
-            registrationIds.forEach(registrationId -> {
+            clients.forEach(client -> {
-                Registration registration = this.clientContext.getRegistration(registrationId);
+                Registration registration = client.getRegistration();
                 Set<String> clientObjects = clientContext.getSupportedIdVerInClient(registration);
                 Set<String> pathSend = analyzerParameters.getPathPostParametersAdd().stream().filter(target -> clientObjects.contains("/" + target.split(LWM2M_SEPARATOR_PATH)[1]))
                         .collect(Collectors.toUnmodifiableSet());
@@ -1160,8 +1166,8 @@ public class DefaultLwM2MTransportMsgHandler implements LwM2mTransportMsgHandler
         }
         // #6.2
         if (analyzerParameters.getPathPostParametersDel().size() > 0) {
-            registrationIds.forEach(registrationId -> {
+            clients.forEach(client -> {
-                Registration registration = this.clientContext.getRegistration(registrationId);
+                Registration registration = client.getRegistration();
                 Set<String> clientObjects = clientContext.getSupportedIdVerInClient(registration);
                 Set<String> pathSend = analyzerParameters.getPathPostParametersDel().stream().filter(target -> clientObjects.contains("/" + target.split(LWM2M_SEPARATOR_PATH)[1]))
                         .collect(Collectors.toUnmodifiableSet());
@@ -1180,7 +1186,7 @@ public class DefaultLwM2MTransportMsgHandler implements LwM2mTransportMsgHandler
     }
 
     private void cancelObserveFromProfile(Registration registration, Set<String> paramAnallyzer) {
-        LwM2mClient lwM2MClient = clientContext.getOrRegister(registration);
+        LwM2mClient lwM2MClient = clientContext.getClientByEndpoint(registration.getEndpoint());
         paramAnallyzer.forEach(pathIdVer -> {
                     if (this.getResourceValueFromLwM2MClient(lwM2MClient, pathIdVer) != null) {
                         lwM2mTransportRequest.sendAllRequest(registration, pathIdVer, OBSERVE_CANCEL, null,
@@ -1199,7 +1205,7 @@ public class DefaultLwM2MTransportMsgHandler implements LwM2mTransportMsgHandler
             log.error("Failed update resource [{}] [{}]", path, valueNew);
             String logMsg = String.format("%s: Failed update resource path - %s value - %s. Value is not changed or bad",
                     LOG_LW2M_ERROR, path, valueNew);
-            this.sendLogsToThingsboard(logMsg, lwM2MClient.getRegistration().getId());
+            this.sendLogsToThingsboard(lwM2MClient, logMsg);
             log.info("Failed update resource [{}] [{}]", path, valueNew);
         }
     }
@@ -1275,8 +1281,7 @@ public class DefaultLwM2MTransportMsgHandler implements LwM2mTransportMsgHandler
                 this.updateResourcesValueToClient(lwM2MClient, this.getResourceValueFormatKv(lwM2MClient, pathIdVer),
                         getValueFromKvProto(tsKvProto.getKv()), pathIdVer);
             });
-        }
+        } else {
-        else {
             log.error("UpdateAttributeFromThingsboard, lwM2MClient is null");
         }
     }
@@ -1285,14 +1290,9 @@ public class DefaultLwM2MTransportMsgHandler implements LwM2mTransportMsgHandler
      * @param lwM2MClient -
      * @return SessionInfoProto -
      */
-    private SessionInfoProto getSessionInfoOrCloseSession(LwM2mClient lwM2MClient) {
+    private SessionInfoProto getSessionInfo(LwM2mClient lwM2MClient) {
-        if (lwM2MClient != null) {
+        if (lwM2MClient != null && lwM2MClient.getSession() != null) {
-            SessionInfoProto sessionInfoProto = lwM2MClient.getSession();
+            return lwM2MClient.getSession();
-            if (sessionInfoProto == null) {
-                log.info("[{}] [{}]", lwM2MClient.getEndpoint(), CLIENT_NOT_AUTHORIZED);
-                this.closeClientSession(lwM2MClient.getRegistration());
-            }
-            return sessionInfoProto;
         }
         return null;
     }
@@ -1302,15 +1302,7 @@ public class DefaultLwM2MTransportMsgHandler implements LwM2mTransportMsgHandler
      * @return - sessionInfo after access connect client
      */
     public SessionInfoProto getSessionInfoOrCloseSession(Registration registration) {
-        return getSessionInfoOrCloseSession(clientContext.getOrRegister(registration));
+        return getSessionInfo(clientContext.getClientByEndpoint(registration.getEndpoint()));
-    }
-
-    /**
-     * @param registrationId -
-     * @return -
-     */
-    private SessionInfoProto getSessionInfoOrCloseSession(String registrationId) {
-        return getSessionInfoOrCloseSession(clientContext.getClientByRegistrationId(registrationId));
     }
 
     /**
@@ -1340,7 +1332,7 @@ public class DefaultLwM2MTransportMsgHandler implements LwM2mTransportMsgHandler
      * @param lwM2MClient - LwM2M Client
      */
     public void putDelayedUpdateResourcesThingsboard(LwM2mClient lwM2MClient) {
-        SessionInfoProto sessionInfo = this.getSessionInfoOrCloseSession(lwM2MClient);
+        SessionInfoProto sessionInfo = this.getSessionInfo(lwM2MClient);
         if (sessionInfo != null) {
             //#1.1
             ConcurrentMap<String, String> keyNamesMap = this.getNamesFromProfileForSharedAttributes(lwM2MClient);
@@ -1359,7 +1351,7 @@ public class DefaultLwM2MTransportMsgHandler implements LwM2mTransportMsgHandler
 
     public void getInfoFirmwareUpdate(LwM2mClient lwM2MClient, Lwm2mClientRpcRequest rpcRequest) {
         if (lwM2MClient.getRegistration().getSupportedVersion(FW_ID) != null) {
-            SessionInfoProto sessionInfo = this.getSessionInfoOrCloseSession(lwM2MClient);
+            SessionInfoProto sessionInfo = this.getSessionInfo(lwM2MClient);
             if (sessionInfo != null) {
                 DefaultLwM2MTransportMsgHandler handler = this;
                 this.transportService.process(sessionInfo, createOtaPackageRequestMsg(sessionInfo, OtaPackageType.FIRMWARE.name()),
@@ -1368,16 +1360,15 @@ public class DefaultLwM2MTransportMsgHandler implements LwM2mTransportMsgHandler
                             public void onSuccess(TransportProtos.GetOtaPackageResponseMsg response) {
                                 if (TransportProtos.ResponseStatus.SUCCESS.equals(response.getResponseStatus())
                                         && response.getType().equals(OtaPackageType.FIRMWARE.name())) {
-                                    log.warn ("7) firmware start with ver: [{}]", response.getVersion());
+                                    log.warn("7) firmware start with ver: [{}]", response.getVersion());
                                     lwM2MClient.getFwUpdate().setRpcRequest(rpcRequest);
                                     lwM2MClient.getFwUpdate().setCurrentVersion(response.getVersion());
                                     lwM2MClient.getFwUpdate().setCurrentTitle(response.getTitle());
                                     lwM2MClient.getFwUpdate().setCurrentId(new OtaPackageId(new UUID(response.getOtaPackageIdMSB(), response.getOtaPackageIdLSB())).getId());
                                     if (rpcRequest == null) {
                                         lwM2MClient.getFwUpdate().sendReadObserveInfo(lwM2mTransportRequest);
-                                    }
+                                    } else {
-                                    else {
+                                        lwM2MClient.getFwUpdate().writeFwSwWare(handler, lwM2mTransportRequest);
-                                        lwM2MClient.getFwUpdate().writeFwSwWare(handler,  lwM2mTransportRequest);
                                     }
                                 } else {
                                     log.trace("OtaPackage [{}] [{}]", lwM2MClient.getDeviceName(), response.getResponseStatus().toString());
@@ -1395,7 +1386,7 @@ public class DefaultLwM2MTransportMsgHandler implements LwM2mTransportMsgHandler
 
     public void getInfoSoftwareUpdate(LwM2mClient lwM2MClient, Lwm2mClientRpcRequest rpcRequest) {
         if (lwM2MClient.getRegistration().getSupportedVersion(SW_ID) != null) {
-            SessionInfoProto sessionInfo = this.getSessionInfoOrCloseSession(lwM2MClient);
+            SessionInfoProto sessionInfo = this.getSessionInfo(lwM2MClient);
             if (sessionInfo != null) {
                 DefaultLwM2MTransportMsgHandler handler = this;
                 transportService.process(sessionInfo, createOtaPackageRequestMsg(sessionInfo, OtaPackageType.SOFTWARE.name()),
@@ -1411,9 +1402,8 @@ public class DefaultLwM2MTransportMsgHandler implements LwM2mTransportMsgHandler
                                     lwM2MClient.getSwUpdate().sendReadObserveInfo(lwM2mTransportRequest);
                                     if (rpcRequest == null) {
                                         lwM2MClient.getSwUpdate().sendReadObserveInfo(lwM2mTransportRequest);
-                                    }
+                                    } else {
-                                    else {
+                                        lwM2MClient.getSwUpdate().writeFwSwWare(handler, lwM2mTransportRequest);
-                                        lwM2MClient.getSwUpdate().writeFwSwWare(handler,  lwM2mTransportRequest);
                                     }
                                 } else {
                                     log.trace("Software [{}] [{}]", lwM2MClient.getDeviceName(), response.getResponseStatus().toString());

common/transport/lwm2m/src/main/java/org/thingsboard/server/transport/lwm2m/server/LwM2mServerListener.java

@@ -87,7 +87,7 @@ public class LwM2mServerListener {
         @Override
         public void cancelled(Observation observation) {
             String msg = String.format("%s:  Canceled Observation  %s.", LOG_LW2M_INFO, observation.getPath());
-            service.sendLogsToThingsboard(msg, observation.getRegistrationId());
+            service.sendLogsToThingsboard2(observation.getRegistrationId(), msg);
             log.warn(msg);
         }
 
@@ -109,7 +109,7 @@ public class LwM2mServerListener {
             String msg = String.format("%s: Successful start newObservation  %s.", LOG_LW2M_INFO,
                     observation.getPath());
             log.warn(msg);
-            service.sendLogsToThingsboard(msg, registration.getId());
+            service.sendLogsToThingsboard2(registration.getId(), msg);
         }
     };
 }

common/transport/lwm2m/src/main/java/org/thingsboard/server/transport/lwm2m/server/LwM2mTransportMsgHandler.java

@@ -22,6 +22,7 @@ import org.thingsboard.server.common.data.Device;
 import org.thingsboard.server.common.data.DeviceProfile;
 import org.thingsboard.server.gen.transport.TransportProtos;
 import org.thingsboard.server.transport.lwm2m.config.LwM2MTransportServerConfig;
+import org.thingsboard.server.transport.lwm2m.server.client.LwM2mClient;
 import org.thingsboard.server.transport.lwm2m.server.client.Lwm2mClientRpcRequest;
 
 import java.util.Collection;
@@ -63,7 +64,9 @@ public interface LwM2mTransportMsgHandler {
 
     void onAwakeDev(Registration registration);
 
-    void sendLogsToThingsboard(String msg, String registrationId);
+    void sendLogsToThingsboard(LwM2mClient client, String msg);
+
+    void sendLogsToThingsboard2(String registrationId, String msg);
 
     LwM2MTransportServerConfig getConfig();
 }

common/transport/lwm2m/src/main/java/org/thingsboard/server/transport/lwm2m/server/LwM2mTransportRequest.java

@@ -127,10 +127,10 @@ public class LwM2mTransportRequest {
 
     public void sendAllRequest(Registration registration, String targetIdVer, LwM2mTypeOper typeOper,
                                String contentFormatName, Object params, long timeoutInMs, Lwm2mClientRpcRequest lwm2mClientRpcRequest) {
+        LwM2mClient lwM2MClient = this.lwM2mClientContext.getClientByEndpoint(registration.getEndpoint());
         try {
             String target = convertPathFromIdVerToObjectId(targetIdVer);
             ContentFormat contentFormat = contentFormatName != null ? ContentFormat.fromName(contentFormatName.toUpperCase()) : ContentFormat.DEFAULT;
-            LwM2mClient lwM2MClient = this.lwM2mClientContext.getOrRegister(registration);
             LwM2mPath resultIds = target != null ? new LwM2mPath(target) : null;
             if (!OBSERVE_CANCEL.name().equals(typeOper.name()) && resultIds != null && registration != null && resultIds.getObjectId() >= 0 && lwM2MClient != null) {
                 if (lwM2MClient.isValidObjectVersion(targetIdVer)) {
@@ -185,7 +185,7 @@ public class LwM2mTransportRequest {
                         }
                         String msg = String.format("%s: type operation %s paths - %s", LOG_LW2M_INFO,
                                 typeOper.name(), paths);
-                        this.handler.sendLogsToThingsboard(msg, registration.getId());
+                        this.handler.sendLogsToThingsboard(lwM2MClient, msg);
                         if (lwm2mClientRpcRequest != null) {
                             String valueMsg = String.format("Paths - %s", paths);
                             this.handler.sentRpcResponse(lwm2mClientRpcRequest, CONTENT.name(), valueMsg, LOG_LW2M_VALUE);
@@ -204,7 +204,7 @@ public class LwM2mTransportRequest {
                             observeCancelMsg = String.format("%s: type operation %s paths: All  count: %d", LOG_LW2M_INFO,
                                     OBSERVE_CANCEL.name(), observeCancelCnt);
                         }
-                        this.afterObserveCancel(registration, observeCancelCnt, observeCancelMsg, lwm2mClientRpcRequest);
+                        this.afterObserveCancel(lwM2MClient, observeCancelCnt, observeCancelMsg, lwm2mClientRpcRequest);
                         break;
                     // lwm2mClientRpcRequest != null
                     case FW_UPDATE:
@@ -215,7 +215,7 @@ public class LwM2mTransportRequest {
         } catch (Exception e) {
             String msg = String.format("%s: type operation %s  %s", LOG_LW2M_ERROR,
                     typeOper.name(), e.getMessage());
-            handler.sendLogsToThingsboard(msg, registration.getId());
+            handler.sendLogsToThingsboard(lwM2MClient, msg);
             if (lwm2mClientRpcRequest != null) {
                 String errorMsg = String.format("Path %s type operation %s  %s", targetIdVer, typeOper.name(), e.getMessage());
                 handler.sentRpcResponse(lwm2mClientRpcRequest, NOT_FOUND.getName(), errorMsg, LOG_LW2M_ERROR);
@@ -273,7 +273,7 @@ public class LwM2mTransportRequest {
                     contentFormat = getContentFormatByResourceModelType(resourceModelWrite, contentFormat);
                     request = this.getWriteRequestSingleResource(contentFormat, resultIds.getObjectId(),
                             resultIds.getObjectInstanceId(), resultIds.getResourceId(), params, resourceModelWrite.type,
-                            registration, rpcRequest);
+                            lwM2MClient, rpcRequest);
                 }
                 break;
             case WRITE_UPDATE:
@@ -337,11 +337,11 @@ public class LwM2mTransportRequest {
                 lwM2MClient.initReadValue(this.handler, convertPathFromObjectIdToIdVer(request.getPath().toString(), registration));
             }
             if (CoAP.ResponseCode.isSuccess(((Response) response.getCoapResponse()).getCode())) {
-                this.handleResponse(registration, request.getPath().toString(), response, request, rpcRequest);
+                this.handleResponse(lwM2MClient, request.getPath().toString(), response, request, rpcRequest);
             } else {
                 String msg = String.format("%s: SendRequest %s: CoapCode - %s Lwm2m code - %d name - %s Resource path - %s", LOG_LW2M_ERROR, request.getClass().getName().toString(),
                         ((Response) response.getCoapResponse()).getCode(), response.getCode().getCode(), response.getCode().getName(), request.getPath().toString());
-                handler.sendLogsToThingsboard(msg, registration.getId());
+                handler.sendLogsToThingsboard(lwM2MClient, msg);
                 log.error("[{}] [{}], [{}] - [{}] [{}] error SendRequest", request.getClass().getName().toString(), registration.getEndpoint(),
                         ((Response) response.getCoapResponse()).getCode(), response.getCode(), request.getPath().toString());
                 if (!lwM2MClient.isInit()) {
@@ -388,7 +388,7 @@ public class LwM2mTransportRequest {
             }
             String msg = String.format("%s: SendRequest %s: Resource path - %s msg error - %s",
                     LOG_LW2M_ERROR, request.getClass().getName().toString(), request.getPath().toString(), e.getMessage());
-            handler.sendLogsToThingsboard(msg, registration.getId());
+            handler.sendLogsToThingsboard(lwM2MClient, msg);
             log.error("[{}] [{}] - [{}] error SendRequest", request.getClass().getName().toString(), request.getPath().toString(), e.toString());
             if (rpcRequest != null) {
                 handler.sentRpcResponse(rpcRequest, CoAP.CodeClass.ERROR_RESPONSE.name(), e.getMessage(), LOG_LW2M_ERROR);
@@ -398,7 +398,7 @@ public class LwM2mTransportRequest {
 
     private WriteRequest getWriteRequestSingleResource(ContentFormat contentFormat, Integer objectId, Integer instanceId,
                                                        Integer resourceId, Object value, ResourceModel.Type type,
-                                                       Registration registration, Lwm2mClientRpcRequest rpcRequest) {
+                                                       LwM2mClient client, Lwm2mClientRpcRequest rpcRequest) {
         try {
             if (type != null) {
                 switch (type) {
@@ -433,7 +433,7 @@ public class LwM2mTransportRequest {
             String patn = "/" + objectId + "/" + instanceId + "/" + resourceId;
             String msg = String.format(LOG_LW2M_ERROR + ": NumberFormatException: Resource path - %s type - %s value - %s msg error - %s  SendRequest to Client",
                     patn, type, value, e.toString());
-            handler.sendLogsToThingsboard(msg, registration.getId());
+            handler.sendLogsToThingsboard(client, msg);
             log.error("Path: [{}] type: [{}] value: [{}] errorMsg: [{}]]", patn, type, value, e.toString());
             if (rpcRequest != null) {
                 String errorMsg = String.format("NumberFormatException: Resource path - %s type - %s value - %s", patn, type, value);
@@ -443,13 +443,13 @@ public class LwM2mTransportRequest {
         }
     }
 
-    private void handleResponse(Registration registration, final String path, LwM2mResponse response,
+    private void handleResponse(LwM2mClient lwM2mClient, final String path, LwM2mResponse response,
                                 SimpleDownlinkRequest request, Lwm2mClientRpcRequest rpcRequest) {
         responseRequestExecutor.submit(() -> {
             try {
-                this.sendResponse(registration, path, response, request, rpcRequest);
+                this.sendResponse(lwM2mClient, path, response, request, rpcRequest);
             } catch (Exception e) {
-                log.error("[{}] endpoint [{}] path [{}] Exception Unable to after send response.", registration.getEndpoint(), path, e);
+                log.error("[{}] endpoint [{}] path [{}] Exception Unable to after send response.", lwM2mClient.getRegistration().getEndpoint(), path, e);
             }
         });
     }
@@ -461,8 +461,9 @@ public class LwM2mTransportRequest {
      * @param path         -
      * @param response     -
      */
-    private void sendResponse(Registration registration, String path, LwM2mResponse response,
+    private void sendResponse(LwM2mClient lwM2mClient, String path, LwM2mResponse response,
                               SimpleDownlinkRequest request, Lwm2mClientRpcRequest rpcRequest) {
+        Registration registration = lwM2mClient.getRegistration();
         String pathIdVer = convertPathFromObjectIdToIdVer(path, registration);
         String msgLog = "";
         if (response instanceof ReadResponse) {
@@ -477,7 +478,7 @@ public class LwM2mTransportRequest {
             String discoverValue = Link.serialize(((DiscoverResponse) response).getObjectLinks());
             msgLog = String.format("%s: type operation: %s path: %s value: %s",
                     LOG_LW2M_INFO, DISCOVER.name(), request.getPath().toString(), discoverValue);
-            handler.sendLogsToThingsboard(msgLog, registration.getId());
+            handler.sendLogsToThingsboard(lwM2mClient, msgLog);
             log.warn("DiscoverResponse: [{}]", (DiscoverResponse) response);
             if (rpcRequest != null) {
                 handler.sentRpcResponse(rpcRequest, response.getCode().getName(), discoverValue, LOG_LW2M_VALUE);
@@ -486,7 +487,7 @@ public class LwM2mTransportRequest {
             msgLog = String.format("%s: type operation: %s path: %s",
                     LOG_LW2M_INFO, EXECUTE.name(), request.getPath().toString());
             log.warn("9) [{}] ", msgLog);
-            handler.sendLogsToThingsboard(msgLog, registration.getId());
+            handler.sendLogsToThingsboard(lwM2mClient, msgLog);
             if (rpcRequest != null) {
                 msgLog = String.format("Start %s path: %S. Preparation finished: %s", EXECUTE.name(), path, rpcRequest.getInfoMsg());
                 rpcRequest.setInfoMsg(msgLog);
@@ -496,7 +497,7 @@ public class LwM2mTransportRequest {
         } else if (response instanceof WriteAttributesResponse) {
             msgLog = String.format("%s: type operation: %s path: %s value: %s",
                     LOG_LW2M_INFO, WRITE_ATTRIBUTES.name(), request.getPath().toString(), ((WriteAttributesRequest) request).getAttributes().toString());
-            handler.sendLogsToThingsboard(msgLog, registration.getId());
+            handler.sendLogsToThingsboard(lwM2mClient, msgLog);
             log.warn("12) [{}] Path [{}] WriteAttributesResponse", pathIdVer, response);
             if (rpcRequest != null) {
                 handler.sentRpcResponse(rpcRequest, response.getCode().getName(), response.toString(), LOG_LW2M_VALUE);
@@ -504,13 +505,14 @@ public class LwM2mTransportRequest {
         } else if (response instanceof WriteResponse) {
             msgLog = String.format("Type operation: Write path: %s", pathIdVer);
             log.warn("10) [{}] response: [{}]", msgLog, response);
-            this.infoWriteResponse(registration, response, request, rpcRequest);
+            this.infoWriteResponse(lwM2mClient, response, request, rpcRequest);
             handler.onWriteResponseOk(registration, pathIdVer, (WriteRequest) request);
         }
     }
 
-    private void infoWriteResponse(Registration registration, LwM2mResponse response, SimpleDownlinkRequest request, Lwm2mClientRpcRequest rpcRequest) {
+    private void infoWriteResponse(LwM2mClient lwM2mClient, LwM2mResponse response, SimpleDownlinkRequest request, Lwm2mClientRpcRequest rpcRequest) {
         try {
+            Registration registration = lwM2mClient.getRegistration();
             LwM2mNode node = ((WriteRequest) request).getNode();
             String msg = null;
             Object value;
@@ -545,7 +547,7 @@ public class LwM2mTransportRequest {
                 }
             }
             if (msg != null) {
-                handler.sendLogsToThingsboard(msg, registration.getId());
+                handler.sendLogsToThingsboard(lwM2mClient, msg);
                 if (request.getPath().toString().equals(FW_PACKAGE_ID) || request.getPath().toString().equals(SW_PACKAGE_ID)) {
                     this.afterWriteSuccessFwSwUpdate(registration, request);
                     if (rpcRequest != null) {
@@ -603,8 +605,8 @@ public class LwM2mTransportRequest {
         }
     }
 
-    private void afterObserveCancel(Registration registration, int observeCancelCnt, String observeCancelMsg, Lwm2mClientRpcRequest rpcRequest) {
+    private void afterObserveCancel(LwM2mClient lwM2mClient, int observeCancelCnt, String observeCancelMsg, Lwm2mClientRpcRequest rpcRequest) {
-        handler.sendLogsToThingsboard(observeCancelMsg, registration.getId());
+        handler.sendLogsToThingsboard(lwM2mClient, observeCancelMsg);
         log.warn("[{}]", observeCancelMsg);
         if (rpcRequest != null) {
             rpcRequest.setInfoMsg(String.format("Count: %d", observeCancelCnt));

common/transport/lwm2m/src/main/java/org/thingsboard/server/transport/lwm2m/server/client/LwM2MClientState.java

@@ -0,0 +1,22 @@
+/**
+ * Copyright © 2016-2021 The Thingsboard Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.thingsboard.server.transport.lwm2m.server.client;
+
+public enum LwM2MClientState {
+
+    CREATED, REGISTERED, UNREGISTERED
+
+}

common/transport/lwm2m/src/main/java/org/thingsboard/server/transport/lwm2m/server/client/LwM2MClientStateException.java

@@ -0,0 +1,31 @@
+/**
+ * Copyright © 2016-2021 The Thingsboard Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.thingsboard.server.transport.lwm2m.server.client;
+
+import lombok.Getter;
+
+public class LwM2MClientStateException extends Exception {
+
+    private static final long serialVersionUID = 3307690997951364046L;
+
+    @Getter
+    private final LwM2MClientState state;
+
+    public LwM2MClientStateException(LwM2MClientState state, String message) {
+        super(message);
+        this.state = state;
+    }
+}

common/transport/lwm2m/src/main/java/org/thingsboard/server/transport/lwm2m/server/client/LwM2mClient.java

@@ -49,6 +49,8 @@ import java.util.UUID;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReentrantLock;
 import java.util.stream.Collectors;
 
 import static org.eclipse.leshan.core.model.ResourceModel.Type.OPAQUE;
@@ -62,12 +64,28 @@ import static org.thingsboard.server.transport.lwm2m.server.LwM2mTransportUtil.g
 
 @Slf4j
 public class LwM2mClient implements Cloneable {
+
+    private final String nodeId;
+    @Getter
+    private final String endpoint;
+    private final Lock lock;
+    @Getter @Setter
+    private LwM2MClientState state;
+    @Getter
+    private final Map<String, ResourceValue> resources;
+    @Getter
+    private final Map<String, TsKvProto> delayedRequests;
+    @Getter
+    @Setter
+    private final List<String> pendingReadRequests;
+    @Getter
+    private final Queue<LwM2mQueuedRequest> queuedRequests;
+
     @Getter
     private String deviceName;
     @Getter
     private String deviceProfileName;
-    @Getter
+
-    private String endpoint;
     @Getter
     private String identity;
     @Getter
@@ -92,15 +110,6 @@ public class LwM2mClient implements Cloneable {
 
     private ValidateDeviceCredentialsResponse credentials;
 
-    @Getter
-    private final Map<String, ResourceValue> resources;
-    @Getter
-    private final Map<String, TsKvProto> delayedRequests;
-    @Getter
-    @Setter
-    private final List<String> pendingReadRequests;
-    @Getter
-    private final Queue<LwM2mQueuedRequest> queuedRequests;
     @Getter
     private boolean init;
 
@@ -108,18 +117,23 @@ public class LwM2mClient implements Cloneable {
         return super.clone();
     }
 
-    public LwM2mClient(String nodeId, String endpoint, String identity, SecurityInfo securityInfo, ValidateDeviceCredentialsResponse credentials, UUID profileId, UUID sessionId) {
+    public LwM2mClient(String nodeId, String endpoint) {
+        this.nodeId = nodeId;
         this.endpoint = endpoint;
-        this.identity = identity;
+        this.lock = new ReentrantLock();
-        this.securityInfo = securityInfo;
-        this.credentials = credentials;
         this.delayedRequests = new ConcurrentHashMap<>();
         this.pendingReadRequests = new CopyOnWriteArrayList<>();
         this.resources = new ConcurrentHashMap<>();
-        this.profileId = profileId;
-        this.init = false;
         this.queuedRequests = new ConcurrentLinkedQueue<>();
+        this.state = LwM2MClientState.CREATED;
+    }
 
+    public void init(String identity, SecurityInfo securityInfo, ValidateDeviceCredentialsResponse credentials, UUID profileId, UUID sessionId) {
+        this.identity = identity;
+        this.securityInfo = securityInfo;
+        this.credentials = credentials;
+        this.profileId = profileId;
+        this.init = false;
         this.fwUpdate = new LwM2mFwSwUpdate(this, OtaPackageType.FIRMWARE);
         this.swUpdate = new LwM2mFwSwUpdate(this, OtaPackageType.SOFTWARE);
         if (this.credentials != null && this.credentials.hasDeviceInfo()) {
@@ -131,6 +145,14 @@ public class LwM2mClient implements Cloneable {
         }
     }
 
+    public void lock() {
+        lock.lock();
+    }
+
+    public void unlock() {
+        lock.unlock();
+    }
+
     public void onDeviceUpdate(Device device, Optional<DeviceProfile> deviceProfileOpt) {
         SessionInfoProto.Builder builder = SessionInfoProto.newBuilder().mergeFrom(session);
         this.deviceId = device.getUuidId();

common/transport/lwm2m/src/main/java/org/thingsboard/server/transport/lwm2m/server/client/LwM2mClientContext.java

@@ -16,6 +16,7 @@
 package org.thingsboard.server.transport.lwm2m.server.client;
 
 import org.eclipse.leshan.server.registration.Registration;
+import org.eclipse.leshan.server.security.SecurityInfo;
 import org.thingsboard.server.common.data.DeviceProfile;
 import org.thingsboard.server.common.transport.auth.ValidateDeviceCredentialsResponse;
 import org.thingsboard.server.gen.transport.TransportProtos;
@@ -27,21 +28,27 @@ import java.util.UUID;
 
 public interface LwM2mClientContext {
 
-    void removeClientByRegistrationId(String registrationId);
+    LwM2mClient getClientByRegistrationId(String registrationId);
 
     LwM2mClient getClientByEndpoint(String endpoint);
 
-    LwM2mClient getClientByRegistrationId(String registrationId);
+    void register(LwM2mClient lwM2MClient, Registration registration) throws LwM2MClientStateException;
+
+    void updateRegistration(LwM2mClient client, Registration registration) throws LwM2MClientStateException;
+
+    void unregister(LwM2mClient client, Registration registration) throws LwM2MClientStateException;
+
+    SecurityInfo fetchSecurityInfoByCredentials(String credentialsId);
 
     LwM2mClient getClient(TransportProtos.SessionInfoProto sessionInfo);
 
-    LwM2mClient getOrRegister(Registration registration);
+//    LwM2mClient getOrRegister(Registration registration);
 
-    LwM2mClient registerOrUpdate(Registration registration);
+//    LwM2mClient registerOrUpdate(Registration registration);
 
-    LwM2mClient fetchClientByEndpoint(String endpoint);
+//    LwM2mClient fetchClientByEndpoint(String endpoint);
 
-    Registration getRegistration(String registrationId);
+//    Registration getRegistration(String registrationId);
 
     Collection<LwM2mClient> getLwM2mClients();
 
@@ -60,4 +67,6 @@ public interface LwM2mClientContext {
     LwM2mClient getClientByDeviceId(UUID deviceId);
 
     void registerClient(Registration registration, ValidateDeviceCredentialsResponse credentials);
+
+
 }

common/transport/lwm2m/src/main/java/org/thingsboard/server/transport/lwm2m/server/client/LwM2mClientContextImpl.java

@@ -27,6 +27,7 @@ import org.thingsboard.server.gen.transport.TransportProtos;
 import org.thingsboard.server.queue.util.TbLwM2mTransportComponent;
 import org.thingsboard.server.transport.lwm2m.secure.EndpointSecurityInfo;
 import org.thingsboard.server.transport.lwm2m.secure.LwM2mCredentialsSecurityInfoValidator;
+import org.thingsboard.server.transport.lwm2m.server.LwM2mQueuedRequest;
 import org.thingsboard.server.transport.lwm2m.server.LwM2mTransportContext;
 import org.thingsboard.server.transport.lwm2m.server.LwM2mTransportUtil;
 
@@ -39,6 +40,8 @@ import java.util.UUID;
 import java.util.concurrent.ConcurrentHashMap;
 
 import static org.eclipse.leshan.core.SecurityMode.NO_SEC;
+import static org.thingsboard.server.transport.lwm2m.server.LwM2mTransportUtil.LOG_LW2M_ERROR;
+import static org.thingsboard.server.transport.lwm2m.server.LwM2mTransportUtil.LOG_LW2M_INFO;
 import static org.thingsboard.server.transport.lwm2m.server.LwM2mTransportUtil.convertPathFromObjectIdToIdVer;
 
 @Slf4j
@@ -58,36 +61,106 @@ public class LwM2mClientContextImpl implements LwM2mClientContext {
 
     @Override
     public LwM2mClient getClientByEndpoint(String endpoint) {
-        return lwM2mClientsByEndpoint.get(endpoint);
+        return lwM2mClientsByEndpoint.computeIfAbsent(endpoint, ep -> new LwM2mClient(context.getNodeId(), ep));
     }
 
     @Override
-    public LwM2mClient getClientByRegistrationId(String registrationId) {
+    public void register(LwM2mClient lwM2MClient, Registration registration) throws LwM2MClientStateException {
-        return lwM2mClientsByRegistrationId.get(registrationId);
+        lwM2MClient.lock();
+        try {
+            if (LwM2MClientState.UNREGISTERED.equals(lwM2MClient.getState())) {
+                throw new LwM2MClientStateException(lwM2MClient.getState(), "Client is in invalid state.");
+            }
+            //TODO: Move this security info lookup to the TbLwM2mSecurityStore.
+            EndpointSecurityInfo securityInfo = lwM2MCredentialsSecurityInfoValidator.getEndpointSecurityInfo(lwM2MClient.getEndpoint(), LwM2mTransportUtil.LwM2mTypeServer.CLIENT);
+            if (securityInfo.getSecurityMode() != null) {
+                if (securityInfo.getDeviceProfile() != null) {
+                    UUID profileUuid = profileUpdate(securityInfo.getDeviceProfile()) != null ? securityInfo.getDeviceProfile().getUuidId() : null;
+                    if (securityInfo.getSecurityInfo() != null) {
+                        lwM2MClient.init(securityInfo.getSecurityInfo().getIdentity(), securityInfo.getSecurityInfo(), securityInfo.getMsg(), profileUuid, UUID.randomUUID());
+                    } else if (NO_SEC.equals(securityInfo.getSecurityMode())) {
+                        lwM2MClient.init(null, null, securityInfo.getMsg(), profileUuid, UUID.randomUUID());
+                    } else {
+                        throw new RuntimeException(String.format("Registration failed: device %s not found.", lwM2MClient.getEndpoint()));
+                    }
+                } else {
+                    throw new RuntimeException(String.format("Registration failed: device %s not found.", lwM2MClient.getEndpoint()));
+                }
+            } else {
+                throw new RuntimeException(String.format("Registration failed: FORBIDDEN, endpointId: %s", lwM2MClient.getEndpoint()));
+            }
+            lwM2MClient.setRegistration(registration);
+            this.lwM2mClientsByRegistrationId.put(registration.getId(), lwM2MClient);
+            lwM2MClient.setState(LwM2MClientState.REGISTERED);
+        } finally {
+            lwM2MClient.unlock();
+        }
     }
 
     @Override
-    public LwM2mClient getOrRegister(Registration registration) {
+    public void updateRegistration(LwM2mClient lwM2MClient, Registration registration) throws LwM2MClientStateException {
-        if (registration == null) {
+        lwM2MClient.lock();
-            return null;
+        try {
+            if (!LwM2MClientState.REGISTERED.equals(lwM2MClient.getState())) {
+                throw new LwM2MClientStateException(lwM2MClient.getState(), "Client is in invalid state.");
+            }
+            Registration currentRegistration = lwM2MClient.getRegistration();
+            if (currentRegistration.getId().equals(registration.getId())) {
+                lwM2MClient.setRegistration(registration);
+            } else {
+                throw new LwM2MClientStateException(lwM2MClient.getState(), "Client has different registration.");
+            }
+        } finally {
+            lwM2MClient.unlock();
         }
-        LwM2mClient client = lwM2mClientsByRegistrationId.get(registration.getId());
+    }
-        if (client == null) {
+
-            client = lwM2mClientsByEndpoint.get(registration.getEndpoint());
+    @Override
-            if (client == null) {
+    public void unregister(LwM2mClient lwM2MClient, Registration registration) throws LwM2MClientStateException {
-                client = registerOrUpdate(registration);
+        lwM2MClient.lock();
+        try {
+            if (!LwM2MClientState.REGISTERED.equals(lwM2MClient.getState())) {
+                throw new LwM2MClientStateException(lwM2MClient.getState(), "Client is in invalid state.");
+            }
+            lwM2mClientsByRegistrationId.remove(registration.getId());
+            Registration currentRegistration = lwM2MClient.getRegistration();
+            if (currentRegistration.getId().equals(registration.getId())) {
+                lwM2MClient.setState(LwM2MClientState.UNREGISTERED);
+                lwM2mClientsByEndpoint.remove(lwM2MClient.getEndpoint());
+                this.securityStore.remove(lwM2MClient.getEndpoint(), false);
+                this.lwM2mClientsByRegistrationId.remove(registration.getId());
+                UUID profileId = lwM2MClient.getProfileId();
+                if (profileId != null) {
+                    Optional<LwM2mClient> otherClients = lwM2mClientsByRegistrationId.values().stream().filter(e -> e.getProfileId().equals(profileId)).findFirst();
+                    if (otherClients.isEmpty()) {
+                        profiles.remove(profileId);
+                    }
+                }
+            } else {
+                throw new LwM2MClientStateException(lwM2MClient.getState(), "Client has different registration.");
             }
+        } finally {
+            lwM2MClient.unlock();
         }
-        return client;
+    }
+
+    @Override
+    public LwM2mClient fetchSecurityInfoByCredentials(String credentialsId) {
+        return null;
+    }
+
+    @Override
+    public LwM2mClient getClientByRegistrationId(String registrationId) {
+        return lwM2mClientsByRegistrationId.get(registrationId);
     }
 
     @Override
     public LwM2mClient getClient(TransportProtos.SessionInfoProto sessionInfo) {
-        LwM2mClient lwM2mClient =  lwM2mClientsByEndpoint.values().stream().filter(c ->
+        LwM2mClient lwM2mClient = lwM2mClientsByEndpoint.values().stream().filter(c ->
                 (new UUID(sessionInfo.getSessionIdMSB(), sessionInfo.getSessionIdLSB()))
                         .equals((new UUID(c.getSession().getSessionIdMSB(), c.getSession().getSessionIdLSB())))
 
-        ).findAny().get();
+        ).findAny().orElse(null);
         if (lwM2mClient == null) {
             log.warn("Device TimeOut? lwM2mClient is null.");
             log.warn("SessionInfo input [{}], lwM2mClientsByEndpoint size: [{}]", sessionInfo, lwM2mClientsByEndpoint.values().size());
@@ -96,60 +169,14 @@ public class LwM2mClientContextImpl implements LwM2mClientContext {
         return lwM2mClient;
     }
 
-    @Override
-    public LwM2mClient registerOrUpdate(Registration registration) {
-        LwM2mClient lwM2MClient = lwM2mClientsByEndpoint.get(registration.getEndpoint());
-        if (lwM2MClient == null) {
-            lwM2MClient = this.fetchClientByEndpoint(registration.getEndpoint());
-        }
-        lwM2MClient.setRegistration(registration);
-//        TODO: this remove is probably redundant. We should remove it.
-//        this.lwM2mClientsByEndpoint.remove(registration.getEndpoint());
-        this.lwM2mClientsByRegistrationId.put(registration.getId(), lwM2MClient);
-        return lwM2MClient;
-    }
-
     public Registration getRegistration(String registrationId) {
         return this.lwM2mClientsByRegistrationId.get(registrationId).getRegistration();
     }
 
-    @Override
-    public LwM2mClient fetchClientByEndpoint(String endpoint) {
-        EndpointSecurityInfo securityInfo = lwM2MCredentialsSecurityInfoValidator.getEndpointSecurityInfo(endpoint, LwM2mTransportUtil.LwM2mTypeServer.CLIENT);
-        if (securityInfo.getSecurityMode() != null) {
-            if (securityInfo.getDeviceProfile() != null) {
-                UUID profileUuid = profileUpdate(securityInfo.getDeviceProfile())!= null ?
-                        securityInfo.getDeviceProfile().getUuidId() : null;
-                //        TODO: for tests bug.
-                if (profileUuid== null) {
-                    log.trace("input parameters toClientProfile if the result is null: [{}]", securityInfo.getDeviceProfile());
-                }
-                LwM2mClient client;
-                if (securityInfo.getSecurityInfo() != null) {
-                    client = new LwM2mClient(context.getNodeId(), securityInfo.getSecurityInfo().getEndpoint(),
-                            securityInfo.getSecurityInfo().getIdentity(), securityInfo.getSecurityInfo(),
-                            securityInfo.getMsg(), profileUuid, UUID.randomUUID());
-                } else if (NO_SEC.equals(securityInfo.getSecurityMode())) {
-                    client = new LwM2mClient(context.getNodeId(), endpoint,
-                            null, null,
-                            securityInfo.getMsg(), profileUuid, UUID.randomUUID());
-                } else {
-                    throw new RuntimeException(String.format("Registration failed: device %s not found.", endpoint));
-                }
-                lwM2mClientsByEndpoint.put(client.getEndpoint(), client);
-                return client;
-            } else {
-                throw new RuntimeException(String.format("Registration failed: device %s not found.", endpoint));
-            }
-        } else {
-            throw new RuntimeException(String.format("Registration failed: FORBIDDEN, endpointId: %s", endpoint));
-        }
-    }
-
     @Override
     public void registerClient(Registration registration, ValidateDeviceCredentialsResponse credentials) {
-        LwM2mClient client = new LwM2mClient(context.getNodeId(), registration.getEndpoint(), null, null, credentials, credentials.getDeviceProfile().getUuidId(), UUID.randomUUID());
+        LwM2mClient client = getClientByEndpoint(registration.getEndpoint());
-        lwM2mClientsByEndpoint.put(registration.getEndpoint(), client);
+        client.init(null, null, credentials, credentials.getDeviceProfile().getUuidId(), UUID.randomUUID());
         lwM2mClientsByRegistrationId.put(registration.getId(), client);
         profileUpdate(credentials.getDeviceProfile());
     }
@@ -171,7 +198,7 @@ public class LwM2mClientContextImpl implements LwM2mClientContext {
 
     @Override
     public LwM2mClientProfile getProfile(Registration registration) {
-        return this.getProfiles().get(getOrRegister(registration).getProfileId());
+        return this.getProfiles().get(getClientByEndpoint(registration.getEndpoint()).getProfileId());
     }
 
     @Override
@@ -186,8 +213,7 @@ public class LwM2mClientContextImpl implements LwM2mClientContext {
         if (lwM2MClientProfile != null) {
             profiles.put(deviceProfile.getUuidId(), lwM2MClientProfile);
             return lwM2MClientProfile;
-        }
+        } else {
-        else {
             return null;
         }
     }
@@ -215,20 +241,4 @@ public class LwM2mClientContextImpl implements LwM2mClientContext {
         return lwM2mClientsByRegistrationId.values().stream().filter(e -> deviceId.equals(e.getDeviceId())).findFirst().orElse(null);
     }
 
-    @Override
-    public void removeClientByRegistrationId(String registrationId) {
-        LwM2mClient lwM2MClient = this.lwM2mClientsByRegistrationId.get(registrationId);
-        if (lwM2MClient != null) {
-            this.securityStore.remove(lwM2MClient.getEndpoint(), false);
-            this.lwM2mClientsByEndpoint.remove(lwM2MClient.getEndpoint());
-            this.lwM2mClientsByRegistrationId.remove(registrationId);
-            UUID profileId = lwM2MClient.getProfileId();
-            if (profileId != null) {
-                Optional<LwM2mClient> otherClients = lwM2mClientsByRegistrationId.values().stream().filter(e -> e.getProfileId().equals(profileId)).findFirst();
-                if (otherClients.isEmpty()) {
-                    profiles.remove(profileId);
-                }
-            }
-        }
-    }
 }

common/transport/lwm2m/src/main/java/org/thingsboard/server/transport/lwm2m/server/client/LwM2mFwSwUpdate.java

@@ -167,7 +167,7 @@ public class LwM2mFwSwUpdate {
             String targetIdVer = convertPathFromObjectIdToIdVer(this.pathPackageId, this.lwM2MClient.getRegistration());
             String fwMsg = String.format("%s: Start type operation %s paths:  %s", LOG_LW2M_INFO,
                     LwM2mTransportUtil.LwM2mTypeOper.FW_UPDATE.name(),  FW_PACKAGE_ID);
-            handler.sendLogsToThingsboard(fwMsg, lwM2MClient.getRegistration().getId());
+            handler.sendLogsToThingsboard(lwM2MClient, fwMsg);
             log.warn("8) Start firmware Update. Send save to: [{}] ver: [{}] path: [{}]", this.lwM2MClient.getDeviceName(), this.currentVersion, targetIdVer);
             request.sendAllRequest(this.lwM2MClient.getRegistration(), targetIdVer, WRITE_REPLACE, ContentFormat.OPAQUE.getName(),
                     firmwareChunk, handler.config.getTimeout(), this.rpcRequest);
@@ -190,7 +190,7 @@ public class LwM2mFwSwUpdate {
         if (LOG_LW2M_ERROR.equals(typeInfo)) {
             msg = String.format("%s Error: %s", msg, msgError);
         }
-        handler.sendLogsToThingsboard(msg, lwM2MClient.getRegistration().getId());
+        handler.sendLogsToThingsboard(lwM2MClient, msg);
     }
 
 

common/transport/lwm2m/src/main/java/org/thingsboard/server/transport/lwm2m/server/store/TbLwM2mSecurityStore.java

@@ -68,7 +68,7 @@ public class TbLwM2mSecurityStore implements EditableSecurityStore {
             if (lwM2mClient != null && lwM2mClient.getRegistration() != null && !lwM2mClient.getRegistration().getIdentity().isSecure()) {
                 return null;
             }
-            securityInfo = clientContext.fetchClientByEndpoint(endpoint).getSecurityInfo();
+            securityInfo = clientContext.fetchSecurityInfoByCredentials(endpoint);
             try {
                 if (securityInfo != null) {
                     add(securityInfo);
@@ -84,7 +84,7 @@ public class TbLwM2mSecurityStore implements EditableSecurityStore {
     public SecurityInfo getByIdentity(String pskIdentity) {
         SecurityInfo securityInfo = securityStore.getByIdentity(pskIdentity);
         if (securityInfo == null) {
-            securityInfo = clientContext.fetchClientByEndpoint(pskIdentity).getSecurityInfo();
+            securityInfo = clientContext.fetchSecurityInfoByCredentials(pskIdentity);
             try {
                 if (securityInfo != null) {
                     add(securityInfo);