From 30d133bc6b9f942efa06bb4e2f5b4d535b1ee253 Mon Sep 17 00:00:00 2001
From: Victor Antonovich <v.antonovich@gmail.com>
Date: Fri, 24 Jun 2016 16:45:26 +0300
Subject: [PATCH] Fix memory leak in BasicCommandStore

---
 .../command/ClusteredExecutionContext.java    | 22 +++++++++----------
 .../cellar/core/command/ResultHandler.java    |  2 +-
 .../cellar/core/command/TimeoutTask.java      |  7 ++----
 3 files changed, 14 insertions(+), 17 deletions(-)

diff --git a/core/src/main/java/org/apache/karaf/cellar/core/command/ClusteredExecutionContext.java b/core/src/main/java/org/apache/karaf/cellar/core/command/ClusteredExecutionContext.java
index 2dbc049b4..a7ee2d06f 100644
--- a/core/src/main/java/org/apache/karaf/cellar/core/command/ClusteredExecutionContext.java
+++ b/core/src/main/java/org/apache/karaf/cellar/core/command/ClusteredExecutionContext.java
@@ -43,21 +43,21 @@ public ClusteredExecutionContext(Producer producer, CommandStore commandStore) {
     }
 
     @Override
-    public <R extends Result, C extends Command<R>> Map<Node, R> execute(C command) throws StoreNotFoundException, ProducerNotFoundException, InterruptedException {
-        if (command == null) {
+    public <R extends Result, C extends Command<R>> Map<Node, R> execute(C command)
+            throws StoreNotFoundException, ProducerNotFoundException, InterruptedException {
+        if (commandStore == null) {
             throw new StoreNotFoundException("Command store not found");
-        } else {
-            commandStore.getPending().put(command.getId(), command);
-            TimeoutTask timeoutTask = new TimeoutTask(command, commandStore);
-            timeoutScheduler.schedule(timeoutTask, command.getTimeout(), TimeUnit.MILLISECONDS);
         }
-
-        if (producer != null) {
-            producer.produce(command);
-            return command.getResult();
-        } else {
+        if (producer == null) {
             throw new ProducerNotFoundException("Command producer not found");
         }
+
+        commandStore.getPending().put(command.getId(), command);
+        TimeoutTask timeoutTask = new TimeoutTask(command, commandStore);
+        timeoutScheduler.schedule(timeoutTask, command.getTimeout(), TimeUnit.MILLISECONDS);
+
+        producer.produce(command);
+        return command.getResult();
     }
 
     public Producer getProducer() {
diff --git a/core/src/main/java/org/apache/karaf/cellar/core/command/ResultHandler.java b/core/src/main/java/org/apache/karaf/cellar/core/command/ResultHandler.java
index 8786e038c..e50aeb30a 100644
--- a/core/src/main/java/org/apache/karaf/cellar/core/command/ResultHandler.java
+++ b/core/src/main/java/org/apache/karaf/cellar/core/command/ResultHandler.java
@@ -37,7 +37,7 @@ public class ResultHandler<R extends Result> implements EventHandler<R> {
     public void handle(R result) {
         if (commandStore != null && commandStore.getPending() != null) {
             String id = result.getId();
-            Command command = commandStore.getPending().get(id);
+            Command command = commandStore.getPending().remove(id);
 
             if (command != null && handlerSwitch.getStatus().equals(SwitchStatus.ON)) {
                 command.addResults(result);
diff --git a/core/src/main/java/org/apache/karaf/cellar/core/command/TimeoutTask.java b/core/src/main/java/org/apache/karaf/cellar/core/command/TimeoutTask.java
index 01f37dc6e..3ac30f400 100644
--- a/core/src/main/java/org/apache/karaf/cellar/core/command/TimeoutTask.java
+++ b/core/src/main/java/org/apache/karaf/cellar/core/command/TimeoutTask.java
@@ -31,11 +31,8 @@ public TimeoutTask(Command command, CommandStore store) {
      */
     @Override
     public void run() {
-        // check if command is still pending
-        Boolean pending = store.getPending().containsKey(command.getId());
-        if (pending) {
-            store.getPending().remove(command.getId());
-        }
+        // remove command if it is still pending
+        store.getPending().remove(command.getId());
     }
 
 }