diff --git a/hadoop-hdds/container-service/src/main/java/org/apache/hadoop/ozone/container/common/statemachine/StateContext.java b/hadoop-hdds/container-service/src/main/java/org/apache/hadoop/ozone/container/common/statemachine/StateContext.java
index 712b3f0f2326..93a45905975a 100644
--- a/hadoop-hdds/container-service/src/main/java/org/apache/hadoop/ozone/container/common/statemachine/StateContext.java
+++ b/hadoop-hdds/container-service/src/main/java/org/apache/hadoop/ozone/container/common/statemachine/StateContext.java
@@ -152,6 +152,8 @@ public class StateContext {
 
   private final String threadNamePrefix;
 
+  private RunningDatanodeState runningDatanodeState;
+
   /**
    * Constructs a StateContext.
    *
@@ -612,9 +614,11 @@ public DatanodeState<DatanodeStateMachine.DatanodeStates> getTask() {
           parentDatanodeStateMachine.getConnectionManager(),
           this);
     case RUNNING:
-      return new RunningDatanodeState(this.conf,
-          parentDatanodeStateMachine.getConnectionManager(),
-          this);
+      if (runningDatanodeState == null) {
+        runningDatanodeState = new RunningDatanodeState(this.conf,
+            parentDatanodeStateMachine.getConnectionManager(), this);
+      }
+      return runningDatanodeState;
     case SHUTDOWN:
       return null;
     default:
@@ -654,7 +658,11 @@ public void execute(ExecutorService service, long time, TimeUnit unit)
     // Adding not null check, in a case where datanode is still starting up, but
     // we called stop DatanodeStateMachine, this sets state to SHUTDOWN, and
     // there is a chance of getting task as null.
-    if (task != null) {
+    if (task == null) {
+      return;
+    }
+
+    try {
       if (this.isEntering()) {
         task.onEnter();
       }
@@ -691,6 +699,8 @@ public void execute(ExecutorService service, long time, TimeUnit unit)
         // that we can terminate the datanode.
         setShutdownOnError();
       }
+    } finally {
+      task.clear();
     }
   }
 
diff --git a/hadoop-hdds/container-service/src/main/java/org/apache/hadoop/ozone/container/common/states/DatanodeState.java b/hadoop-hdds/container-service/src/main/java/org/apache/hadoop/ozone/container/common/states/DatanodeState.java
index 25be207dcd9c..f057a852bc4e 100644
--- a/hadoop-hdds/container-service/src/main/java/org/apache/hadoop/ozone/container/common/states/DatanodeState.java
+++ b/hadoop-hdds/container-service/src/main/java/org/apache/hadoop/ozone/container/common/states/DatanodeState.java
@@ -55,4 +55,10 @@ public interface DatanodeState<T> {
   T await(long time, TimeUnit timeUnit)
       throws InterruptedException, ExecutionException, TimeoutException;
 
+  /**
+   * Clean up some resources.
+   */
+  @SuppressWarnings("checkstyle:WhitespaceAround")
+  default void clear(){}
+
 }
diff --git a/hadoop-hdds/container-service/src/main/java/org/apache/hadoop/ozone/container/common/states/datanode/RunningDatanodeState.java b/hadoop-hdds/container-service/src/main/java/org/apache/hadoop/ozone/container/common/states/datanode/RunningDatanodeState.java
index 4e5b64c27e5e..c4d507668d4c 100644
--- a/hadoop-hdds/container-service/src/main/java/org/apache/hadoop/ozone/container/common/states/datanode/RunningDatanodeState.java
+++ b/hadoop-hdds/container-service/src/main/java/org/apache/hadoop/ozone/container/common/states/datanode/RunningDatanodeState.java
@@ -31,11 +31,8 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import java.util.EnumMap;
-import java.util.HashMap;
 import java.util.LinkedList;
 import java.util.List;
-import java.util.Map;
 import java.util.concurrent.Callable;
 import java.util.concurrent.CompletionService;
 import java.util.concurrent.ExecutionException;
@@ -54,9 +51,6 @@ public class RunningDatanodeState implements DatanodeState {
   private final ConfigurationSource conf;
   private final StateContext context;
   private CompletionService<EndPointStates> ecs;
-  /** Cache the end point task per end point per end point state. */
-  private Map<EndpointStateMachine, Map<EndPointStates,
-      Callable<EndPointStates>>> endpointTasks;
 
   public RunningDatanodeState(ConfigurationSource conf,
       SCMConnectionManager connectionManager,
@@ -64,55 +58,6 @@ public RunningDatanodeState(ConfigurationSource conf,
     this.connectionManager = connectionManager;
     this.conf = conf;
     this.context = context;
-    initEndPointTask();
-  }
-
-  /**
-   * Initialize end point tasks corresponding to each end point,
-   * each end point state.
-   */
-  private void initEndPointTask() {
-    endpointTasks = new HashMap<>();
-    for (EndpointStateMachine endpoint : connectionManager.getValues()) {
-      EnumMap<EndPointStates, Callable<EndPointStates>> endpointTaskForState =
-          new EnumMap<>(EndPointStates.class);
-
-      for (EndPointStates state : EndPointStates.values()) {
-        Callable<EndPointStates> endPointTask = null;
-        switch (state) {
-        case GETVERSION:
-          endPointTask = new VersionEndpointTask(endpoint, conf,
-              context.getParent().getContainer());
-          break;
-        case REGISTER:
-          endPointTask = RegisterEndpointTask.newBuilder()
-              .setConfig(conf)
-              .setEndpointStateMachine(endpoint)
-              .setContext(context)
-              .setDatanodeDetails(context.getParent().getDatanodeDetails())
-              .setOzoneContainer(context.getParent().getContainer())
-              .build();
-          break;
-        case HEARTBEAT:
-          endPointTask = HeartbeatEndpointTask.newBuilder()
-              .setConfig(conf)
-              .setEndpointStateMachine(endpoint)
-              .setDatanodeDetails(context
-                  .getParent()
-                  .getDatanodeDetails())
-              .setContext(context)
-              .build();
-          break;
-        default:
-          break;
-        }
-
-        if (endPointTask != null) {
-          endpointTaskForState.put(state, endPointTask);
-        }
-      }
-      endpointTasks.put(endpoint, endpointTaskForState);
-    }
   }
 
   /**
@@ -140,7 +85,7 @@ public void onExit() {
   public void execute(ExecutorService executor) {
     ecs = new ExecutorCompletionService<>(executor);
     for (EndpointStateMachine endpoint : connectionManager.getValues()) {
-      Callable<EndPointStates> endpointTask = getEndPointTask(endpoint);
+      Callable<EndPointStates> endpointTask = buildEndPointTask(endpoint);
       if (endpointTask != null) {
         // Just do a timely wait. A slow EndpointStateMachine won't occupy
         // the thread in executor from DatanodeStateMachine for a long time,
@@ -171,12 +116,30 @@ public void setExecutorCompletionService(ExecutorCompletionService e) {
     this.ecs = e;
   }
 
-  private Callable<EndPointStates> getEndPointTask(
+  @SuppressWarnings("checkstyle:Indentation")
+  private Callable<EndPointStates> buildEndPointTask(
       EndpointStateMachine endpoint) {
-    if (endpointTasks.containsKey(endpoint)) {
-      return endpointTasks.get(endpoint).get(endpoint.getState());
-    } else {
-      throw new IllegalArgumentException("Illegal endpoint: " + endpoint);
+    switch (endpoint.getState()) {
+      case GETVERSION:
+        return new VersionEndpointTask(endpoint, conf,
+            context.getParent().getContainer());
+      case REGISTER:
+        return RegisterEndpointTask.newBuilder()
+            .setConfig(conf)
+            .setEndpointStateMachine(endpoint)
+            .setContext(context)
+            .setDatanodeDetails(context.getParent().getDatanodeDetails())
+            .setOzoneContainer(context.getParent().getContainer())
+            .build();
+      case HEARTBEAT:
+        return HeartbeatEndpointTask.newBuilder()
+            .setConfig(conf)
+            .setEndpointStateMachine(endpoint)
+            .setDatanodeDetails(context.getParent().getDatanodeDetails())
+            .setContext(context)
+            .build();
+      default:
+        return null;
     }
   }
 
@@ -238,4 +201,9 @@ private Callable<EndPointStates> getEndPointTask(
     }
     return computeNextContainerState(results);
   }
+
+  @Override
+  public void clear() {
+    ecs = null;
+  }
 }