001/*
002 * Licensed to the Apache Software Foundation (ASF) under one or more
003 * contributor license agreements.  See the NOTICE file distributed with
004 * this work for additional information regarding copyright ownership.
005 * The ASF licenses this file to You under the Apache License, Version 2.0
006 * (the "License"); you may not use this file except in compliance with
007 * the License.  You may obtain a copy of the License at
008 *
009 *      http://www.apache.org/licenses/LICENSE-2.0
010 *
011 * Unless required by applicable law or agreed to in writing, software
012 * distributed under the License is distributed on an "AS IS" BASIS,
013 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
014 * See the License for the specific language governing permissions and
015 * limitations under the License.
016 */
017package org.apache.activemq.store.memory;
018
019import java.io.IOException;
020import java.util.Collections;
021import java.util.LinkedHashMap;
022import java.util.Map;
023
024import org.apache.activemq.broker.ConnectionContext;
025import org.apache.activemq.command.ActiveMQDestination;
026import org.apache.activemq.command.Message;
027import org.apache.activemq.command.MessageAck;
028import org.apache.activemq.command.MessageId;
029import org.apache.activemq.store.AbstractMessageStore;
030import org.apache.activemq.store.IndexListener;
031import org.apache.activemq.store.MessageRecoveryListener;
032import org.apache.activemq.store.MessageStoreStatistics;
033
034/**
035 * An implementation of {@link org.apache.activemq.store.MessageStore}
036 */
037public class MemoryMessageStore extends AbstractMessageStore {
038
039    protected final Map<MessageId, Message> messageTable;
040    protected MessageId lastBatchId;
041    protected long sequenceId;
042
043    public MemoryMessageStore(ActiveMQDestination destination) {
044        this(destination, new LinkedHashMap<MessageId, Message>());
045    }
046
047    public MemoryMessageStore(ActiveMQDestination destination, Map<MessageId, Message> messageTable) {
048        super(destination);
049        this.messageTable = Collections.synchronizedMap(messageTable);
050    }
051
052    @Override
053    public synchronized void addMessage(ConnectionContext context, Message message) throws IOException {
054        synchronized (messageTable) {
055            messageTable.put(message.getMessageId(), message);
056            incMessageStoreStatistics(getMessageStoreStatistics(), message);
057            message.incrementReferenceCount();
058            message.getMessageId().setFutureOrSequenceLong(sequenceId++);
059            if (indexListener != null) {
060                indexListener.onAdd(new IndexListener.MessageContext(context, message, null));
061            }
062        }
063    }
064
065    @Override
066    public Message getMessage(MessageId identity) throws IOException {
067        return messageTable.get(identity);
068    }
069
070    @Override
071    public void removeMessage(ConnectionContext context, MessageAck ack) throws IOException {
072        removeMessage(ack.getLastMessageId());
073    }
074
075    public void removeMessage(MessageId msgId) throws IOException {
076        synchronized (messageTable) {
077            Message removed = messageTable.remove(msgId);
078            if (removed != null) {
079                removed.decrementReferenceCount();
080                decMessageStoreStatistics(getMessageStoreStatistics(), removed);
081            }
082            if ((lastBatchId != null && lastBatchId.equals(msgId)) || messageTable.isEmpty()) {
083                lastBatchId = null;
084            }
085        }
086    }
087
088    @Override
089    public void recover(MessageRecoveryListener listener) throws Exception {
090        // the message table is a synchronizedMap - so just have to synchronize here
091        synchronized (messageTable) {
092            for (Message message : messageTable.values()) {
093                listener.recoverMessage(message);
094            }
095        }
096    }
097
098    @Override
099    public void removeAllMessages(ConnectionContext context) throws IOException {
100        synchronized (messageTable) {
101            messageTable.clear();
102            getMessageStoreStatistics().reset();
103        }
104    }
105
106    public void delete() {
107        synchronized (messageTable) {
108            messageTable.clear();
109            getMessageStoreStatistics().reset();
110        }
111    }
112
113    @Override
114    public void recoverNextMessages(int maxReturned, MessageRecoveryListener listener) throws Exception {
115        synchronized (messageTable) {
116            boolean pastLackBatch = lastBatchId == null;
117            for (Map.Entry<MessageId, Message> entry : messageTable.entrySet()) {
118                if (pastLackBatch) {
119                    Object msg = entry.getValue();
120                    lastBatchId = entry.getKey();
121                    if (msg.getClass() == MessageId.class) {
122                        listener.recoverMessageReference((MessageId) msg);
123                    } else {
124                        listener.recoverMessage((Message) msg);
125                    }
126                } else {
127                    pastLackBatch = entry.getKey().equals(lastBatchId);
128                }
129            }
130        }
131    }
132
133    @Override
134    public void resetBatching() {
135        lastBatchId = null;
136    }
137
138    @Override
139    public void setBatch(MessageId messageId) {
140        synchronized (messageTable) {
141            if (messageTable.containsKey(messageId)) {
142                lastBatchId = messageId;
143            } else {
144                resetBatching();
145            }
146        }
147    }
148
149    @Override
150    public void updateMessage(Message message) {
151        synchronized (messageTable) {
152            Message original = messageTable.get(message.getMessageId());
153
154            // if can't be found then increment count, else remove old size
155            if (original == null) {
156                getMessageStoreStatistics().getMessageCount().increment();
157            } else {
158                getMessageStoreStatistics().getMessageSize().addSize(-original.getSize());
159            }
160            messageTable.put(message.getMessageId(), message);
161            getMessageStoreStatistics().getMessageSize().addSize(message.getSize());
162        }
163    }
164
165    @Override
166    public void recoverMessageStoreStatistics() throws IOException {
167        synchronized (messageTable) {
168            long size = 0;
169            int count = 0;
170            for (Message message : messageTable.values()) {
171                size += message.getSize();
172            }
173
174            getMessageStoreStatistics().reset();
175            getMessageStoreStatistics().getMessageCount().setCount(count);
176            getMessageStoreStatistics().getMessageSize().setTotalSize(size);
177        }
178    }
179
180    protected static final void incMessageStoreStatistics(final MessageStoreStatistics stats, final Message message) {
181        if (stats != null && message != null) {
182            stats.getMessageCount().increment();
183            stats.getMessageSize().addSize(message.getSize());
184        }
185    }
186
187    protected static final void decMessageStoreStatistics(final MessageStoreStatistics stats, final Message message) {
188        if (stats != null && message != null) {
189            stats.getMessageCount().decrement();
190            stats.getMessageSize().addSize(-message.getSize());
191        }
192    }
193}