From 4cbbefdd269de1fcf11de622f629ebbf0821b00e Mon Sep 17 00:00:00 2001
From: YevhenBondarenko <ybondarenko@thingsboard.io>
Date: Thu, 29 Jul 2021 14:48:59 +0300
Subject: [PATCH] fw update improvements

---
 .../ota/DefaultLwM2MOtaUpdateService.java     | 34 +++++++++++++++----
 .../lwm2m/server/ota/LwM2MClientOtaInfo.java  |  5 +++
 2 files changed, 32 insertions(+), 7 deletions(-)

diff --git a/common/transport/lwm2m/src/main/java/org/thingsboard/server/transport/lwm2m/server/ota/DefaultLwM2MOtaUpdateService.java b/common/transport/lwm2m/src/main/java/org/thingsboard/server/transport/lwm2m/server/ota/DefaultLwM2MOtaUpdateService.java
index d39cfd94f0..2972e26fba 100644
--- a/common/transport/lwm2m/src/main/java/org/thingsboard/server/transport/lwm2m/server/ota/DefaultLwM2MOtaUpdateService.java
+++ b/common/transport/lwm2m/src/main/java/org/thingsboard/server/transport/lwm2m/server/ota/DefaultLwM2MOtaUpdateService.java
@@ -45,11 +45,11 @@ import org.thingsboard.server.transport.lwm2m.server.downlink.TbLwM2MExecuteRequ
 import org.thingsboard.server.transport.lwm2m.server.downlink.TbLwM2MWriteReplaceRequest;
 import org.thingsboard.server.transport.lwm2m.server.downlink.TbLwM2MWriteResponseCallback;
 import org.thingsboard.server.transport.lwm2m.server.log.LwM2MTelemetryLogService;
-import org.thingsboard.server.transport.lwm2m.server.ota.firmware.LwM2MClientFwOtaInfo;
-import org.thingsboard.server.transport.lwm2m.server.ota.firmware.LwM2MFirmwareUpdateStrategy;
 import org.thingsboard.server.transport.lwm2m.server.ota.firmware.FirmwareDeliveryMethod;
 import org.thingsboard.server.transport.lwm2m.server.ota.firmware.FirmwareUpdateResult;
 import org.thingsboard.server.transport.lwm2m.server.ota.firmware.FirmwareUpdateState;
+import org.thingsboard.server.transport.lwm2m.server.ota.firmware.LwM2MClientFwOtaInfo;
+import org.thingsboard.server.transport.lwm2m.server.ota.firmware.LwM2MFirmwareUpdateStrategy;
 import org.thingsboard.server.transport.lwm2m.server.ota.software.LwM2MClientSwOtaInfo;
 import org.thingsboard.server.transport.lwm2m.server.ota.software.LwM2MSoftwareUpdateStrategy;
 import org.thingsboard.server.transport.lwm2m.server.ota.software.SoftwareUpdateResult;
@@ -178,7 +178,7 @@ public class DefaultLwM2MOtaUpdateService extends LwM2MExecutorAwareService impl
             attributesToFetch.add(SOFTWARE_TAG);
             attributesToFetch.add(SOFTWARE_URL);
         }
-        
+
         var clientSettings = clientContext.getProfile(client.getProfileId()).getClientLwM2mSettings();
 
         onFirmwareStrategyUpdate(client, clientSettings);
@@ -283,8 +283,17 @@ public class DefaultLwM2MOtaUpdateService extends LwM2MExecutorAwareService impl
         }
         fwInfo.setUpdateState(state);
         Optional<OtaPackageUpdateStatus> status = toOtaPackageUpdateStatus(state);
-        status.ifPresent(otaStatus -> sendStateUpdateToTelemetry(client, fwInfo,
-                otaStatus, "Firmware Update State: " + state.name()));
+
+        if (FirmwareUpdateState.IDLE.equals(state) && DOWNLOADING.equals(fwInfo.getStatus())) {
+            fwInfo.setFailedPackageId(fwInfo.getTargetPackageId());
+            status = Optional.of(FAILED);
+        }
+
+        status.ifPresent(otaStatus -> {
+            fwInfo.setStatus(otaStatus);
+            sendStateUpdateToTelemetry(client, fwInfo,
+                    otaStatus, "Firmware Update State: " + state.name());
+        });
         update(fwInfo);
     }
 
@@ -294,8 +303,19 @@ public class DefaultLwM2MOtaUpdateService extends LwM2MExecutorAwareService impl
         LwM2MClientFwOtaInfo fwInfo = getOrInitFwInfo(client);
         FirmwareUpdateResult result = FirmwareUpdateResult.fromUpdateResultFwByCode(code.intValue());
         Optional<OtaPackageUpdateStatus> status = toOtaPackageUpdateStatus(result);
-        status.ifPresent(otaStatus -> sendStateUpdateToTelemetry(client, fwInfo,
-                otaStatus, "Firmware Update Result: " + result.name()));
+
+        if (FirmwareUpdateResult.INITIAL.equals(result) && OtaPackageUpdateStatus.UPDATING.equals(fwInfo.getStatus())) {
+            status = Optional.of(UPDATED);
+            fwInfo.setRetryAttempts(0);
+        }
+
+        status.ifPresent(otaStatus -> {
+                    fwInfo.setStatus(otaStatus);
+                    sendStateUpdateToTelemetry(client, fwInfo,
+                            otaStatus, "Firmware Update Result: " + result.name());
+                }
+        );
+
         if (result.isAgain() && fwInfo.getRetryAttempts() <= 2) {
             fwInfo.setRetryAttempts(fwInfo.getRetryAttempts() + 1);
             startFirmwareUpdateIfNeeded(client, fwInfo);
diff --git a/common/transport/lwm2m/src/main/java/org/thingsboard/server/transport/lwm2m/server/ota/LwM2MClientOtaInfo.java b/common/transport/lwm2m/src/main/java/org/thingsboard/server/transport/lwm2m/server/ota/LwM2MClientOtaInfo.java
index f7aa5521d0..7261e5bd40 100644
--- a/common/transport/lwm2m/src/main/java/org/thingsboard/server/transport/lwm2m/server/ota/LwM2MClientOtaInfo.java
+++ b/common/transport/lwm2m/src/main/java/org/thingsboard/server/transport/lwm2m/server/ota/LwM2MClientOtaInfo.java
@@ -20,6 +20,7 @@ import lombok.Data;
 import lombok.NoArgsConstructor;
 import org.thingsboard.server.common.data.StringUtils;
 import org.thingsboard.server.common.data.ota.OtaPackageType;
+import org.thingsboard.server.common.data.ota.OtaPackageUpdateStatus;
 
 import java.util.Optional;
 
@@ -39,6 +40,7 @@ public abstract class LwM2MClientOtaInfo<Strategy, State, Result> {
     protected Strategy strategy;
     protected State updateState;
     protected Result result;
+    protected OtaPackageUpdateStatus status;
 
     protected String failedPackageId;
     protected int retryAttempts;
@@ -104,4 +106,7 @@ public abstract class LwM2MClientOtaInfo<Strategy, State, Result> {
 
     public abstract OtaPackageType getType();
 
+    public String getTargetPackageId() {
+        return getPackageId(targetName, targetVersion);
+    }
 }