Resolving review comments

provider/notification-azure/src/main/java/org/opengroup/osdu/notification/provider/azure/messageBus/ProcessNotification.java

@@ -32,6 +32,8 @@ import org.slf4j.MDC;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnExpression;
 import org.springframework.stereotype.Component;
+import java.util.concurrent.ConcurrentMap;
+import java.util.Date;
 
 @Component
 @ConditionalOnExpression("${azure.serviceBus.enabled:true} || ${azure.eventGridToServiceBus.enabled:true}")
@@ -50,7 +52,7 @@ public class ProcessNotification {
     public void performNotification(IMessage message, String subscriptionName) throws Exception {
         TelemetryClient telemetryClient = new TelemetryClient();
         StopWatch stopWatch = new StopWatch();
-        stopWatch.start();
+        stopWatch.start();        
 
         try {
             NotificationContent notificationContent = requestBodyAdapter.extractNotificationContent(message, subscriptionName);
@@ -58,13 +60,9 @@ public class ProcessNotification {
             String dataPartitionId = notificationContent.getExtractAttributes().get(DpsHeaders.DATA_PARTITION_ID);
             String correlationId = notificationContent.getExtractAttributes().get(DpsHeaders.CORRELATION_ID);
 
-            MDC.setContextMap(mdcContextMap.getContextMap(correlationId, dataPartitionId));
-            dpsHeaders.setThreadContext(dataPartitionId, correlationId);
-
-            LOGGER.info("Notification process started for message with id: {}", message.getMessageId());
-
-            HttpResponse response = notificationHandler.notifySubscriber(notificationContent.getNotificationId(),
-                    notificationContent.getData(), notificationContent.getExtractAttributes());
+            ConcurrentMap properties = telemetryClient.getContext().getProperties();
+            properties.put("correlation-id",correlationId);
+            properties.put("data-partition-id", dataPartitionId); 
 
             RequestTelemetry requestTelemetry = new RequestTelemetry(
                 "SBQueueRequest",
@@ -72,8 +70,16 @@ public class ProcessNotification {
                 stopWatch.getTime(),
                 "500",
                 false);
-
             requestTelemetry.setId(message.getMessageId());
+            
+            MDC.setContextMap(mdcContextMap.getContextMap(correlationId, dataPartitionId));
+            dpsHeaders.setThreadContext(dataPartitionId, correlationId);
+
+            LOGGER.info("Notification process started for message with id: {}", message.getMessageId());
+
+            HttpResponse response = notificationHandler.notifySubscriber(notificationContent.getNotificationId(),
+                    notificationContent.getData(), notificationContent.getExtractAttributes());
+                                           
             if (!response.isSuccessCode()) {
                 telemetryClient.trackRequest(requestTelemetry); 
                 throw new Exception(NOT_ACKNOWLEDGE);
@@ -83,11 +89,12 @@ public class ProcessNotification {
                 telemetryClient.trackRequest(requestTelemetry);
             }
         } catch (Exception e) {
-            LOGGER.error(String.format("An error occurred performing Notification for message with ID: ", message.getMessageId()), e);
+            LOGGER.error(String.format("An error occurred performing Notification for message with ID: ", message.getMessageId()), e);            
             throw e;
         } finally {
             ThreadScopeContextHolder.getContext().clear();
             MDC.clear();
+            stopWatch.stop();
         }
     }
 }