/*
 * Copyright 2012-2016, the original author or 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 com.flipkart.flux.impl.redriver;

import akka.actor.ActorSystem;
import akka.actor.PoisonPill;
import akka.actor.Props;
import akka.cluster.singleton.ClusterSingletonManager;
import akka.cluster.singleton.ClusterSingletonManagerSettings;
import com.flipkart.flux.client.runtime.FluxRuntimeConnector;
import com.flipkart.flux.impl.boot.ActorSystemManager;
import com.flipkart.flux.redriver.model.ScheduledMessage;
import com.flipkart.flux.redriver.service.MessageManagerService;
import com.flipkart.flux.redriver.service.RedriverService;
import com.flipkart.flux.task.redriver.RedriverRegistry;
import com.flipkart.polyguice.core.Initializable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.slf4j.MDC;

import javax.inject.Inject;
import javax.inject.Singleton;
import static com.flipkart.flux.Constants.TASK_ID;

/**
 * <code>RedriverRegistryImpl</code> is an implementation of the {@link RedriverRegistry} that uses
 * {@link FluxRuntimeConnector} to submit taskExecution request and redriver's {@link MessageManagerService}
 * to register/deregister tasks
 *
 * @author regunath.balasubramanian
 * @author gaurav.ashok
 */
@Singleton
public class RedriverRegistryImpl implements RedriverRegistry, Initializable {

    private static final Logger logger = LoggerFactory.getLogger(RedriverRegistryImpl.class);

    /**
     * ActorSystemManager to create a singleton actor of {@link AkkaRedriverService}
     */
    ActorSystemManager actorSystemManager;

    /**
     * Flux runtime to submit task for execution.
     */
    private FluxRuntimeConnector fluxRuntimeConnector;

    /**
     * Redriver message service to register/deregister tasks
     */
    private MessageManagerService redriverMessageService;

    /**
     * Redriver service responsible for redriving any stuck tasks
     */
    private RedriverService redriverService;

    @Inject
    public RedriverRegistryImpl(ActorSystemManager actorSystemManager, FluxRuntimeConnector fluxRuntimeConnector,
                                MessageManagerService redriverMessageService, RedriverService redriverService) {
        this.actorSystemManager = actorSystemManager;
        this.fluxRuntimeConnector = fluxRuntimeConnector;
        this.redriverMessageService = redriverMessageService;
        this.redriverService = redriverService;
    }

    /**
     * Initialize the singleton actor wrapping redriverService.
     *
     * @see AkkaRedriverService
     */
    @Override
    public void initialize() {
        ActorSystem actorSystem = actorSystemManager.retrieveActorSystem();
        Props actorProps = Props.create(AkkaRedriverService.class, redriverService);
        ClusterSingletonManagerSettings settings = ClusterSingletonManagerSettings.create(actorSystem);
        actorSystem.actorOf(ClusterSingletonManager.props(actorProps, PoisonPill.getInstance(), settings), "redriverServiceActor");
    }

    /**
     * RedriverRegistry method. Registers the Task with the redriver
     *
     * @see RedriverRegistry#redriveTask(String, Long)
     */
    public void registerTask(Long taskId, String stateMachineId, long redriveDelay) {
        logger.info("Register task : {} for redriver with time : {}", taskId, redriveDelay);
        redriverMessageService.saveMessage(new ScheduledMessage(taskId, stateMachineId, System.currentTimeMillis() + redriveDelay));
    }

    /**
     * RedriverRegistry method. Un-Registers the Task with the redriver
     *
     * @see RedriverRegistry#deRegisterTask(String, Long)
     */
    public void deRegisterTask(String stateMachineId, Long taskId) {
        MDC.clear();
        MDC.put(TASK_ID, taskId.toString());
        logger.info("DeRegister task : {} with smId : {} redriver", taskId, stateMachineId);
        redriverMessageService.scheduleForRemoval(stateMachineId, taskId);
    }

    /**
     * RedriverRegistry method. Re-drives the task identified by the <code>taskId</code> by submitting it back to the runtime.
     *
     * @see RedriverRegistry#redriveTask(String, Long)
     */
    public void redriveTask(String stateMachineId, Long taskId) {
        MDC.clear();
        MDC.put(TASK_ID, taskId.toString());
        logger.info("Redrive task with stateMachineId: {} Id : {} ", stateMachineId, taskId);
        fluxRuntimeConnector.redriveTask(stateMachineId, taskId);
    }
}

@@ -1254,16 +1254,61 @@
 izable;%0A
+import com.flipkart.flux.utils.LoggingUtils;%0A
 import o
@@ -1359,30 +1359,8 @@
 ory;
-%0Aimport org.slf4j.MDC;
 %0A%0Aim
@@ -1419,59 +1419,8 @@
 ton;
-%0Aimport static com.flipkart.flux.Constants.TASK_ID;
 %0A%0A/*
@@ -1725,16 +1725,37 @@
 v.ashok%0A
+ * @author akif.khan%0A
  */%0A@Sin
@@ -3732,24 +3732,128 @@
 iveDelay) %7B%0A
+        try %7B%0A            LoggingUtils.registerStateMachineIdForLogging(stateMachineId.toString());%0A    
         logg
@@ -3933,16 +3933,20 @@
 Delay);%0A
+    
         
@@ -4066,24 +4066,117 @@
 iveDelay));%0A
+        %7D finally %7B%0A            LoggingUtils.deRegisterStateMachineIdForLogging();%0A        %7D%0A
     %7D%0A%0A    /
@@ -4404,50 +4404,84 @@
     
-MDC.clear();%0A        MDC.put(TASK_ID, task
+try %7B%0A            LoggingUtils.registerStateMachineIdForLogging(stateMachine
 Id.t
@@ -4492,32 +4492,36 @@
 ing());%0A        
+ 
+   
 logger.info(%22DeR
@@ -4589,32 +4589,36 @@
 ineId);%0A        
+    
 redriverMessageS
@@ -4664,24 +4664,117 @@
 d, taskId);%0A
+        %7D finally %7B%0A            LoggingUtils.deRegisterStateMachineIdForLogging();%0A        %7D%0A
     %7D%0A%0A    /
@@ -5051,50 +5051,84 @@
     
-MDC.clear();%0A        MDC.put(TASK_ID, task
+try %7B%0A            LoggingUtils.registerStateMachineIdForLogging(stateMachine
 Id.t
@@ -5131,32 +5131,36 @@
 Id.toString());%0A
+    
         logger.i
@@ -5245,16 +5245,20 @@
         
+    
 fluxRunt
@@ -5303,20 +5303,113 @@
 d, taskId);%0A
+        %7D finally %7B%0A            LoggingUtils.deRegisterStateMachineIdForLogging();%0A        %7D%0A
     %7D%0A%7D%0A

Use LoggingUtils in RedriverRegistryImpl