From ce1d85d1ee11c669a5f08dcbc4f8514f11cee981 Mon Sep 17 00:00:00 2001
From: "Hiram R. Chirino" <chirino@apache.org>
Date: Thu, 3 Apr 2008 19:26:04 +0000
Subject: [PATCH] Do better reference count management of persistent messages
 when they are stored in a Memory store.

git-svn-id: https://svn.apache.org/repos/asf/activemq/trunk@644447 13f79535-47bb-0310-9956-ffa450edef68
---
 .../region/cursors/AbstractStoreCursor.java   |  4 +++-
 .../store/memory/MemoryMessageStore.java      |  6 +++++-
 .../activemq/store/memory/MemoryTopicSub.java | 21 +++++++++++++------
 3 files changed, 23 insertions(+), 8 deletions(-)

diff --git a/activemq-core/src/main/java/org/apache/activemq/broker/region/cursors/AbstractStoreCursor.java b/activemq-core/src/main/java/org/apache/activemq/broker/region/cursors/AbstractStoreCursor.java
index 0d771caafa..44c85627c8 100644
--- a/activemq-core/src/main/java/org/apache/activemq/broker/region/cursors/AbstractStoreCursor.java
+++ b/activemq-core/src/main/java/org/apache/activemq/broker/region/cursors/AbstractStoreCursor.java
@@ -77,7 +77,9 @@ public abstract class AbstractStoreCursor extends AbstractPendingMessageCursor i
         if (!isDuplicate(message.getMessageId())) {
             if (!cached) {
                 message.setRegionDestination(regionDestination);
-                message.setMemoryUsage(this.getSystemUsage().getMemoryUsage());
+                if( message.getMemoryUsage()==null ) {
+                    message.setMemoryUsage(this.getSystemUsage().getMemoryUsage());
+                }
             }
             message.incrementReferenceCount();
             batchList.put(message.getMessageId(), message);
diff --git a/activemq-core/src/main/java/org/apache/activemq/store/memory/MemoryMessageStore.java b/activemq-core/src/main/java/org/apache/activemq/store/memory/MemoryMessageStore.java
index 059e8678ed..e87567a391 100755
--- a/activemq-core/src/main/java/org/apache/activemq/store/memory/MemoryMessageStore.java
+++ b/activemq-core/src/main/java/org/apache/activemq/store/memory/MemoryMessageStore.java
@@ -58,6 +58,7 @@ public class MemoryMessageStore implements MessageStore {
         synchronized (messageTable) {
             messageTable.put(message.getMessageId(), message);
         }
+        message.incrementReferenceCount();
     }
 
     // public void addMessageReference(ConnectionContext context,MessageId
@@ -82,7 +83,10 @@ public class MemoryMessageStore implements MessageStore {
 
     public void removeMessage(MessageId msgId) throws IOException {
         synchronized (messageTable) {
-            messageTable.remove(msgId);
+            Message removed = messageTable.remove(msgId);
+            if( removed !=null ) {
+                removed.decrementReferenceCount();
+            }
             if ((lastBatchId != null && lastBatchId.equals(msgId)) || messageTable.isEmpty()) {
                 lastBatchId = null;
             }
diff --git a/activemq-core/src/main/java/org/apache/activemq/store/memory/MemoryTopicSub.java b/activemq-core/src/main/java/org/apache/activemq/store/memory/MemoryTopicSub.java
index 13a3a24f8b..c38cea2bf3 100755
--- a/activemq-core/src/main/java/org/apache/activemq/store/memory/MemoryTopicSub.java
+++ b/activemq-core/src/main/java/org/apache/activemq/store/memory/MemoryTopicSub.java
@@ -34,14 +34,23 @@ class MemoryTopicSub {
     private Map<MessageId, Message> map = new LinkedHashMap<MessageId, Message>();
     private MessageId lastBatch;
 
-    synchronized void addMessage(MessageId id, Message message) {
-        map.put(id, message);
+    void addMessage(MessageId id, Message message) {
+        synchronized(this) {
+            map.put(id, message);
+        }
+        message.incrementReferenceCount();
     }
 
-    synchronized void removeMessage(MessageId id) {
-        map.remove(id);
-        if ((lastBatch != null && lastBatch.equals(id)) || map.isEmpty()) {
-            resetBatching();
+    void removeMessage(MessageId id) {
+        Message removed;
+        synchronized(this) {
+            removed = map.remove(id);
+            if ((lastBatch != null && lastBatch.equals(id)) || map.isEmpty()) {
+                resetBatching();
+            }
+        }
+        if( removed!=null ) {
+            removed.decrementReferenceCount();
         }
     }