Source code

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.ArrayList;
021import java.util.Collections;
022import java.util.Iterator;
023import java.util.LinkedHashMap;
024import java.util.List;
025import java.util.Map;
026import java.util.concurrent.ConcurrentHashMap;
027import java.util.concurrent.ConcurrentMap;
028
029import org.apache.activemq.broker.ConnectionContext;
030import org.apache.activemq.command.Message;
031import org.apache.activemq.command.MessageAck;
032import org.apache.activemq.command.MessageId;
033import org.apache.activemq.command.TransactionId;
034import org.apache.activemq.command.XATransactionId;
035import org.apache.activemq.store.InlineListenableFuture;
036import org.apache.activemq.store.ListenableFuture;
037import org.apache.activemq.store.MessageStore;
038import org.apache.activemq.store.PersistenceAdapter;
039import org.apache.activemq.store.ProxyMessageStore;
040import org.apache.activemq.store.ProxyTopicMessageStore;
041import org.apache.activemq.store.TopicMessageStore;
042import org.apache.activemq.store.TransactionRecoveryListener;
043import org.apache.activemq.store.TransactionStore;
044
045/**
046 * Provides a TransactionStore implementation that can create transaction aware
047 * MessageStore objects from non transaction aware MessageStore objects.
048 *
049 *
050 */
051public class MemoryTransactionStore implements TransactionStore {
052
053    protected ConcurrentMap<Object, Tx> inflightTransactions = new ConcurrentHashMap<Object, Tx>();
054    protected Map<TransactionId, Tx> preparedTransactions = Collections.synchronizedMap(new LinkedHashMap<TransactionId, Tx>());
055    protected final PersistenceAdapter persistenceAdapter;
056
057    private boolean doingRecover;
058
059    public class Tx {
060
061        public List<AddMessageCommand> messages = Collections.synchronizedList(new ArrayList<AddMessageCommand>());
062
063        public final List<RemoveMessageCommand> acks = Collections.synchronizedList(new ArrayList<RemoveMessageCommand>());
064
065        public void add(AddMessageCommand msg) {
066            messages.add(msg);
067        }
068
069        public void add(RemoveMessageCommand ack) {
070            acks.add(ack);
071        }
072
073        public Message[] getMessages() {
074            Message rc[] = new Message[messages.size()];
075            int count = 0;
076            for (Iterator<AddMessageCommand> iter = messages.iterator(); iter.hasNext();) {
077                AddMessageCommand cmd = iter.next();
078                rc[count++] = cmd.getMessage();
079            }
080            return rc;
081        }
082
083        public MessageAck[] getAcks() {
084            MessageAck rc[] = new MessageAck[acks.size()];
085            int count = 0;
086            for (Iterator<RemoveMessageCommand> iter = acks.iterator(); iter.hasNext();) {
087                RemoveMessageCommand cmd = iter.next();
088                rc[count++] = cmd.getMessageAck();
089            }
090            return rc;
091        }
092
093        /**
094         * @throws IOException
095         */
096        public void commit() throws IOException {
097            ConnectionContext ctx = new ConnectionContext();
098            persistenceAdapter.beginTransaction(ctx);
099            try {
100
101                // Do all the message adds.
102                for (Iterator<AddMessageCommand> iter = messages.iterator(); iter.hasNext();) {
103                    AddMessageCommand cmd = iter.next();
104                    cmd.run(ctx);
105                }
106                // And removes..
107                for (Iterator<RemoveMessageCommand> iter = acks.iterator(); iter.hasNext();) {
108                    RemoveMessageCommand cmd = iter.next();
109                    cmd.run(ctx);
110                }
111
112            } catch ( IOException e ) {
113                persistenceAdapter.rollbackTransaction(ctx);
114                throw e;
115            }
116            persistenceAdapter.commitTransaction(ctx);
117        }
118    }
119
120    public interface AddMessageCommand {
121        Message getMessage();
122
123        MessageStore getMessageStore();
124
125        void run(ConnectionContext context) throws IOException;
126
127        void setMessageStore(MessageStore messageStore);
128    }
129
130    public interface RemoveMessageCommand {
131        MessageAck getMessageAck();
132
133        void run(ConnectionContext context) throws IOException;
134
135        MessageStore getMessageStore();
136    }
137
138    public MemoryTransactionStore(PersistenceAdapter persistenceAdapter) {
139        this.persistenceAdapter=persistenceAdapter;
140    }
141
142    public MessageStore proxy(MessageStore messageStore) {
143        ProxyMessageStore proxyMessageStore = new ProxyMessageStore(messageStore) {
144            @Override
145            public void addMessage(ConnectionContext context, final Message send) throws IOException {
146                MemoryTransactionStore.this.addMessage(context, getDelegate(), send);
147            }
148
149            @Override
150            public void addMessage(ConnectionContext context, final Message send, boolean canOptimize) throws IOException {
151                MemoryTransactionStore.this.addMessage(context, getDelegate(), send);
152            }
153
154            @Override
155            public ListenableFuture<Object> asyncAddQueueMessage(ConnectionContext context, Message message) throws IOException {
156                MemoryTransactionStore.this.addMessage(context, getDelegate(), message);
157                return new InlineListenableFuture();
158             }
159
160            @Override
161            public ListenableFuture<Object> asyncAddQueueMessage(ConnectionContext context, Message message, boolean canoptimize) throws IOException {
162                MemoryTransactionStore.this.addMessage(context, getDelegate(), message);
163                return new InlineListenableFuture();
164             }
165
166            @Override
167            public void removeMessage(ConnectionContext context, final MessageAck ack) throws IOException {
168                MemoryTransactionStore.this.removeMessage(getDelegate(), ack);
169            }
170
171            @Override
172            public void removeAsyncMessage(ConnectionContext context, MessageAck ack) throws IOException {
173                MemoryTransactionStore.this.removeMessage(getDelegate(), ack);
174            }
175        };
176        onProxyQueueStore(proxyMessageStore);
177        return proxyMessageStore;
178    }
179
180    protected void onProxyQueueStore(ProxyMessageStore proxyMessageStore) {
181    }
182
183    public TopicMessageStore proxy(TopicMessageStore messageStore) {
184        ProxyTopicMessageStore proxyTopicMessageStore = new ProxyTopicMessageStore(messageStore) {
185            @Override
186            public void addMessage(ConnectionContext context, final Message send) throws IOException {
187                MemoryTransactionStore.this.addMessage(context, getDelegate(), send);
188            }
189
190            @Override
191            public void addMessage(ConnectionContext context, final Message send, boolean canOptimize) throws IOException {
192                MemoryTransactionStore.this.addMessage(context, getDelegate(), send);
193            }
194
195            @Override
196            public ListenableFuture<Object> asyncAddTopicMessage(ConnectionContext context, Message message) throws IOException {
197                MemoryTransactionStore.this.addMessage(context, getDelegate(), message);
198                return new InlineListenableFuture();
199             }
200
201            @Override
202            public ListenableFuture<Object> asyncAddTopicMessage(ConnectionContext context, Message message, boolean canOptimize) throws IOException {
203                MemoryTransactionStore.this.addMessage(context, getDelegate(), message);
204                return new InlineListenableFuture();
205             }
206
207            @Override
208            public void removeMessage(ConnectionContext context, final MessageAck ack) throws IOException {
209                MemoryTransactionStore.this.removeMessage(getDelegate(), ack);
210            }
211
212            @Override
213            public void removeAsyncMessage(ConnectionContext context, MessageAck ack) throws IOException {
214                MemoryTransactionStore.this.removeMessage(getDelegate(), ack);
215            }
216
217            @Override
218            public void acknowledge(ConnectionContext context, String clientId, String subscriptionName,
219                            MessageId messageId, MessageAck ack) throws IOException {
220                MemoryTransactionStore.this.acknowledge((TopicMessageStore)getDelegate(), clientId,
221                        subscriptionName, messageId, ack);
222            }
223        };
224        onProxyTopicStore(proxyTopicMessageStore);
225        return proxyTopicMessageStore;
226    }
227
228    protected void onProxyTopicStore(ProxyTopicMessageStore proxyTopicMessageStore) {
229    }
230
231    /**
232     * @see org.apache.activemq.store.TransactionStore#prepare(TransactionId)
233     */
234    @Override
235    public void prepare(TransactionId txid) throws IOException {
236        Tx tx = inflightTransactions.remove(txid);
237        if (tx == null) {
238            return;
239        }
240        preparedTransactions.put(txid, tx);
241    }
242
243    public Tx getTx(Object txid) {
244        Tx tx = inflightTransactions.get(txid);
245        if (tx == null) {
246            synchronized (inflightTransactions) {
247                tx = inflightTransactions.get(txid);
248                if ( tx == null) {
249                    tx = new Tx();
250                    inflightTransactions.put(txid, tx);
251                }
252            }
253        }
254        return tx;
255    }
256
257    public Tx getPreparedTx(TransactionId txid) {
258        Tx tx = preparedTransactions.get(txid);
259        if (tx == null) {
260            tx = new Tx();
261            preparedTransactions.put(txid, tx);
262        }
263        return tx;
264    }
265
266    @Override
267    public void commit(TransactionId txid, boolean wasPrepared, Runnable preCommit,Runnable postCommit) throws IOException {
268        if (preCommit != null) {
269            preCommit.run();
270        }
271        Tx tx;
272        if (wasPrepared) {
273            tx = preparedTransactions.remove(txid);
274        } else {
275            tx = inflightTransactions.remove(txid);
276        }
277
278        if (tx != null) {
279            tx.commit();
280        }
281        if (postCommit != null) {
282            postCommit.run();
283        }
284    }
285
286    /**
287     * @see org.apache.activemq.store.TransactionStore#rollback(TransactionId)
288     */
289    @Override
290    public void rollback(TransactionId txid) throws IOException {
291        preparedTransactions.remove(txid);
292        inflightTransactions.remove(txid);
293    }
294
295    @Override
296    public void start() throws Exception {
297    }
298
299    @Override
300    public void stop() throws Exception {
301    }
302
303    @Override
304    public synchronized void recover(TransactionRecoveryListener listener) throws IOException {
305        // All the inflight transactions get rolled back..
306        inflightTransactions.clear();
307        this.doingRecover = true;
308        try {
309            for (Iterator<TransactionId> iter = preparedTransactions.keySet().iterator(); iter.hasNext();) {
310                Object txid = iter.next();
311                Tx tx = preparedTransactions.get(txid);
312                listener.recover((XATransactionId)txid, tx.getMessages(), tx.getAcks());
313                onRecovered(tx);
314            }
315        } finally {
316            this.doingRecover = false;
317        }
318    }
319
320    protected void onRecovered(Tx tx) {
321    }
322
323    /**
324     * @param message
325     * @throws IOException
326     */
327    void addMessage(final ConnectionContext context, final MessageStore destination, final Message message) throws IOException {
328
329        if (doingRecover) {
330            return;
331        }
332
333        if (message.getTransactionId() != null) {
334            Tx tx = getTx(message.getTransactionId());
335            tx.add(new AddMessageCommand() {
336                MessageStore messageStore = destination;
337                @Override
338                public Message getMessage() {
339                    return message;
340                }
341
342                @Override
343                public MessageStore getMessageStore() {
344                    return destination;
345                }
346
347                @Override
348                public void run(ConnectionContext ctx) throws IOException {
349                    destination.addMessage(ctx, message);
350                }
351
352                @Override
353                public void setMessageStore(MessageStore messageStore) {
354                    this.messageStore = messageStore;
355                }
356
357            });
358        } else {
359            destination.addMessage(context, message);
360        }
361    }
362
363    /**
364     * @param ack
365     * @throws IOException
366     */
367    final void removeMessage(final MessageStore destination, final MessageAck ack) throws IOException {
368        if (doingRecover) {
369            return;
370        }
371
372        if (ack.isInTransaction()) {
373            Tx tx = getTx(ack.getTransactionId());
374            tx.add(new RemoveMessageCommand() {
375                @Override
376                public MessageAck getMessageAck() {
377                    return ack;
378                }
379
380                @Override
381                public void run(ConnectionContext ctx) throws IOException {
382                    destination.removeMessage(ctx, ack);
383                }
384
385                @Override
386                public MessageStore getMessageStore() {
387                    return destination;
388                }
389            });
390        } else {
391            destination.removeMessage(null, ack);
392        }
393    }
394
395    public void acknowledge(final TopicMessageStore destination, final String clientId, final String subscriptionName,
396                           final MessageId messageId, final MessageAck ack) throws IOException {
397        if (doingRecover) {
398            return;
399        }
400
401        if (ack.isInTransaction()) {
402            Tx tx = getTx(ack.getTransactionId());
403            tx.add(new RemoveMessageCommand() {
404                @Override
405                public MessageAck getMessageAck() {
406                    return ack;
407                }
408
409                @Override
410                public void run(ConnectionContext ctx) throws IOException {
411                    destination.acknowledge(ctx, clientId, subscriptionName, messageId, ack);
412                }
413
414                @Override
415                public MessageStore getMessageStore() {
416                    return destination;
417                }
418            });
419        } else {
420            destination.acknowledge(null, clientId, subscriptionName, messageId, ack);
421        }
422    }
423
424
425    public void delete() {
426        inflightTransactions.clear();
427        preparedTransactions.clear();
428        doingRecover = false;
429    }
430
431}