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.kahadb;
018
019import java.io.DataInputStream;
020import java.io.IOException;
021import java.util.ArrayList;
022import java.util.Iterator;
023import java.util.List;
024import java.util.Map;
025import java.util.concurrent.CancellationException;
026import java.util.concurrent.ConcurrentHashMap;
027import java.util.concurrent.ExecutionException;
028import java.util.concurrent.Future;
029
030import org.apache.activemq.broker.ConnectionContext;
031import org.apache.activemq.command.Message;
032import org.apache.activemq.command.MessageAck;
033import org.apache.activemq.command.MessageId;
034import org.apache.activemq.command.TransactionId;
035import org.apache.activemq.command.XATransactionId;
036import org.apache.activemq.protobuf.Buffer;
037import org.apache.activemq.store.AbstractMessageStore;
038import org.apache.activemq.store.ListenableFuture;
039import org.apache.activemq.store.MessageStore;
040import org.apache.activemq.store.ProxyMessageStore;
041import org.apache.activemq.store.ProxyTopicMessageStore;
042import org.apache.activemq.store.TopicMessageStore;
043import org.apache.activemq.store.TransactionRecoveryListener;
044import org.apache.activemq.store.TransactionStore;
045import org.apache.activemq.store.kahadb.MessageDatabase.Operation;
046import org.apache.activemq.store.kahadb.data.KahaCommitCommand;
047import org.apache.activemq.store.kahadb.data.KahaPrepareCommand;
048import org.apache.activemq.store.kahadb.data.KahaRollbackCommand;
049import org.apache.activemq.store.kahadb.data.KahaTransactionInfo;
050import org.apache.activemq.wireformat.WireFormat;
051import org.slf4j.Logger;
052import org.slf4j.LoggerFactory;
053
054/**
055 * Provides a TransactionStore implementation that can create transaction aware
056 * MessageStore objects from non transaction aware MessageStore objects.
057 *
058 *
059 */
060public class KahaDBTransactionStore implements TransactionStore {
061    static final Logger LOG = LoggerFactory.getLogger(KahaDBTransactionStore.class);
062    ConcurrentHashMap<Object, Tx> inflightTransactions = new ConcurrentHashMap<Object, Tx>();
063    private final KahaDBStore theStore;
064
065    public KahaDBTransactionStore(KahaDBStore theStore) {
066        this.theStore = theStore;
067    }
068
069    private WireFormat wireFormat(){
070      return this.theStore.wireFormat;
071    }
072
073    public class Tx {
074        private final ArrayList<AddMessageCommand> messages = new ArrayList<AddMessageCommand>();
075
076        private final ArrayList<RemoveMessageCommand> acks = new ArrayList<RemoveMessageCommand>();
077
078        public void add(AddMessageCommand msg) {
079            messages.add(msg);
080        }
081
082        public void add(RemoveMessageCommand ack) {
083            acks.add(ack);
084        }
085
086        public Message[] getMessages() {
087            Message rc[] = new Message[messages.size()];
088            int count = 0;
089            for (Iterator<AddMessageCommand> iter = messages.iterator(); iter.hasNext();) {
090                AddMessageCommand cmd = iter.next();
091                rc[count++] = cmd.getMessage();
092            }
093            return rc;
094        }
095
096        public MessageAck[] getAcks() {
097            MessageAck rc[] = new MessageAck[acks.size()];
098            int count = 0;
099            for (Iterator<RemoveMessageCommand> iter = acks.iterator(); iter.hasNext();) {
100                RemoveMessageCommand cmd = iter.next();
101                rc[count++] = cmd.getMessageAck();
102            }
103            return rc;
104        }
105
106        /**
107         * @return true if something to commit
108         * @throws IOException
109         */
110        public List<Future<Object>> commit() throws IOException {
111            List<Future<Object>> results = new ArrayList<Future<Object>>();
112            // Do all the message adds.
113            for (Iterator<AddMessageCommand> iter = messages.iterator(); iter.hasNext();) {
114                AddMessageCommand cmd = iter.next();
115                results.add(cmd.run());
116
117            }
118            // And removes..
119            for (Iterator<RemoveMessageCommand> iter = acks.iterator(); iter.hasNext();) {
120                RemoveMessageCommand cmd = iter.next();
121                cmd.run();
122                results.add(cmd.run());
123            }
124
125            return results;
126        }
127    }
128
129    public abstract class AddMessageCommand {
130        private final ConnectionContext ctx;
131        AddMessageCommand(ConnectionContext ctx) {
132            this.ctx = ctx;
133        }
134        abstract Message getMessage();
135        Future<Object> run() throws IOException {
136            return run(this.ctx);
137        }
138        abstract Future<Object> run(ConnectionContext ctx) throws IOException;
139    }
140
141    public abstract class RemoveMessageCommand {
142
143        private final ConnectionContext ctx;
144        RemoveMessageCommand(ConnectionContext ctx) {
145            this.ctx = ctx;
146        }
147        abstract MessageAck getMessageAck();
148        Future<Object> run() throws IOException {
149            return run(this.ctx);
150        }
151        abstract Future<Object> run(ConnectionContext context) throws IOException;
152    }
153
154    public MessageStore proxy(MessageStore messageStore) {
155        return new ProxyMessageStore(messageStore) {
156            @Override
157            public void addMessage(ConnectionContext context, final Message send) throws IOException {
158                KahaDBTransactionStore.this.addMessage(context, getDelegate(), send);
159            }
160
161            @Override
162            public void addMessage(ConnectionContext context, final Message send, boolean canOptimize) throws IOException {
163                KahaDBTransactionStore.this.addMessage(context, getDelegate(), send);
164            }
165
166            @Override
167            public ListenableFuture<Object> asyncAddQueueMessage(ConnectionContext context, Message message) throws IOException {
168                return KahaDBTransactionStore.this.asyncAddQueueMessage(context, getDelegate(), message);
169            }
170
171            @Override
172            public ListenableFuture<Object> asyncAddQueueMessage(ConnectionContext context, Message message, boolean canOptimize) throws IOException {
173                return KahaDBTransactionStore.this.asyncAddQueueMessage(context, getDelegate(), message);
174            }
175
176            @Override
177            public void removeMessage(ConnectionContext context, final MessageAck ack) throws IOException {
178                KahaDBTransactionStore.this.removeMessage(context, getDelegate(), ack);
179            }
180
181            @Override
182            public void removeAsyncMessage(ConnectionContext context, MessageAck ack) throws IOException {
183                KahaDBTransactionStore.this.removeAsyncMessage(context, getDelegate(), ack);
184            }
185        };
186    }
187
188    public TopicMessageStore proxy(TopicMessageStore messageStore) {
189        return new ProxyTopicMessageStore(messageStore) {
190            @Override
191            public void addMessage(ConnectionContext context, final Message send) throws IOException {
192                KahaDBTransactionStore.this.addMessage(context, getDelegate(), send);
193            }
194
195            @Override
196            public void addMessage(ConnectionContext context, final Message send, boolean canOptimize) throws IOException {
197                KahaDBTransactionStore.this.addMessage(context, getDelegate(), send);
198            }
199
200            @Override
201            public ListenableFuture<Object> asyncAddTopicMessage(ConnectionContext context, Message message) throws IOException {
202                return KahaDBTransactionStore.this.asyncAddTopicMessage(context, getDelegate(), message);
203            }
204
205            @Override
206            public ListenableFuture<Object> asyncAddTopicMessage(ConnectionContext context, Message message, boolean canOptimize) throws IOException {
207                return KahaDBTransactionStore.this.asyncAddTopicMessage(context, getDelegate(), message);
208            }
209
210            @Override
211            public void removeMessage(ConnectionContext context, final MessageAck ack) throws IOException {
212                KahaDBTransactionStore.this.removeMessage(context, getDelegate(), ack);
213            }
214
215            @Override
216            public void removeAsyncMessage(ConnectionContext context, MessageAck ack) throws IOException {
217                KahaDBTransactionStore.this.removeAsyncMessage(context, getDelegate(), ack);
218            }
219
220            @Override
221            public void acknowledge(ConnectionContext context, String clientId, String subscriptionName,
222                            MessageId messageId, MessageAck ack) throws IOException {
223                KahaDBTransactionStore.this.acknowledge(context, (TopicMessageStore)getDelegate(), clientId,
224                        subscriptionName, messageId, ack);
225            }
226
227        };
228    }
229
230    /**
231     * @throws IOException
232     * @see org.apache.activemq.store.TransactionStore#prepare(TransactionId)
233     */
234    public void prepare(TransactionId txid) throws IOException {
235        KahaTransactionInfo info = getTransactionInfo(txid);
236        if (txid.isXATransaction() || theStore.isConcurrentStoreAndDispatchTransactions() == false) {
237            theStore.store(new KahaPrepareCommand().setTransactionInfo(info), true, null, null);
238        } else {
239            Tx tx = inflightTransactions.remove(txid);
240            if (tx != null) {
241               theStore.store(new KahaPrepareCommand().setTransactionInfo(info), true, null, null);
242            }
243        }
244    }
245
246    public Tx getTx(Object txid) {
247        Tx tx = inflightTransactions.get(txid);
248        if (tx == null) {
249            tx = new Tx();
250            inflightTransactions.put(txid, tx);
251        }
252        return tx;
253    }
254
255    public void commit(TransactionId txid, boolean wasPrepared, final Runnable preCommit, Runnable postCommit)
256            throws IOException {
257        if (txid != null) {
258            if (!txid.isXATransaction() && theStore.isConcurrentStoreAndDispatchTransactions()) {
259                if (preCommit != null) {
260                    preCommit.run();
261                }
262                Tx tx = inflightTransactions.remove(txid);
263                if (tx != null) {
264                    List<Future<Object>> results = tx.commit();
265                    boolean doneSomething = false;
266                    for (Future<Object> result : results) {
267                        try {
268                            result.get();
269                        } catch (InterruptedException e) {
270                            theStore.brokerService.handleIOException(new IOException(e.getMessage()));
271                        } catch (ExecutionException e) {
272                            theStore.brokerService.handleIOException(new IOException(e.getMessage()));
273                        }catch(CancellationException e) {
274                        }
275                        if (!result.isCancelled()) {
276                            doneSomething = true;
277                        }
278                    }
279                    if (postCommit != null) {
280                        postCommit.run();
281                    }
282                    if (doneSomething) {
283                        KahaTransactionInfo info = getTransactionInfo(txid);
284                        theStore.store(new KahaCommitCommand().setTransactionInfo(info), theStore.isEnableJournalDiskSyncs(), null, null);
285                    }
286                }else {
287                    //The Tx will be null for failed over clients - lets run their post commits
288                    if (postCommit != null) {
289                        postCommit.run();
290                    }
291                }
292
293            } else {
294                KahaTransactionInfo info = getTransactionInfo(txid);
295                if (preCommit != null) {
296                    preCommit.run();
297                }
298                theStore.store(new KahaCommitCommand().setTransactionInfo(info), theStore.isEnableJournalDiskSyncs(), null, postCommit);
299                forgetRecoveredAcks(txid, false);
300            }
301        }else {
302           LOG.error("Null transaction passed on commit");
303        }
304    }
305
306    /**
307     * @throws IOException
308     * @see org.apache.activemq.store.TransactionStore#rollback(TransactionId)
309     */
310    public void rollback(TransactionId txid) throws IOException {
311        if (txid.isXATransaction() || theStore.isConcurrentStoreAndDispatchTransactions() == false) {
312            KahaTransactionInfo info = getTransactionInfo(txid);
313            theStore.store(new KahaRollbackCommand().setTransactionInfo(info), theStore.isEnableJournalDiskSyncs(), null, null);
314            forgetRecoveredAcks(txid, true);
315        } else {
316            inflightTransactions.remove(txid);
317        }
318    }
319
320    protected void forgetRecoveredAcks(TransactionId txid, boolean isRollback) throws IOException {
321        if (txid.isXATransaction()) {
322            XATransactionId xaTid = ((XATransactionId) txid);
323            theStore.forgetRecoveredAcks(xaTid.getPreparedAcks(), isRollback);
324        }
325    }
326
327    public void start() throws Exception {
328    }
329
330    public void stop() throws Exception {
331    }
332
333    public synchronized void recover(TransactionRecoveryListener listener) throws IOException {
334        for (Map.Entry<TransactionId, List<Operation>> entry : theStore.preparedTransactions.entrySet()) {
335            XATransactionId xid = (XATransactionId) entry.getKey();
336            ArrayList<Message> messageList = new ArrayList<Message>();
337            ArrayList<MessageAck> ackList = new ArrayList<MessageAck>();
338
339            for (Operation op : entry.getValue()) {
340                if (op.getClass() == MessageDatabase.AddOperation.class) {
341                    MessageDatabase.AddOperation addOp = (MessageDatabase.AddOperation) op;
342                    Message msg = (Message) wireFormat().unmarshal(new DataInputStream(addOp.getCommand().getMessage()
343                            .newInput()));
344                    messageList.add(msg);
345                } else {
346                    MessageDatabase.RemoveOperation rmOp = (MessageDatabase.RemoveOperation) op;
347                    Buffer ackb = rmOp.getCommand().getAck();
348                    MessageAck ack = (MessageAck) wireFormat().unmarshal(new DataInputStream(ackb.newInput()));
349                    ackList.add(ack);
350                }
351            }
352
353            Message[] addedMessages = new Message[messageList.size()];
354            MessageAck[] acks = new MessageAck[ackList.size()];
355            messageList.toArray(addedMessages);
356            ackList.toArray(acks);
357            xid.setPreparedAcks(ackList);
358            theStore.trackRecoveredAcks(ackList);
359            listener.recover(xid, addedMessages, acks);
360        }
361    }
362
363    /**
364     * @param message
365     * @throws IOException
366     */
367    void addMessage(ConnectionContext context, final MessageStore destination, final Message message)
368            throws IOException {
369
370        if (message.getTransactionId() != null) {
371            if (message.getTransactionId().isXATransaction() || theStore.isConcurrentStoreAndDispatchTransactions() == false) {
372                destination.addMessage(context, message);
373            } else {
374                Tx tx = getTx(message.getTransactionId());
375                tx.add(new AddMessageCommand(context) {
376                    @Override
377                    public Message getMessage() {
378                        return message;
379                    }
380                    @Override
381                    public Future<Object> run(ConnectionContext ctx) throws IOException {
382                        destination.addMessage(ctx, message);
383                        return AbstractMessageStore.FUTURE;
384                    }
385
386                });
387            }
388        } else {
389            destination.addMessage(context, message);
390        }
391    }
392
393    ListenableFuture<Object> asyncAddQueueMessage(ConnectionContext context, final MessageStore destination, final Message message)
394            throws IOException {
395
396        if (message.getTransactionId() != null) {
397            if (message.getTransactionId().isXATransaction() || theStore.isConcurrentStoreAndDispatchTransactions() == false) {
398                destination.addMessage(context, message);
399                return AbstractMessageStore.FUTURE;
400            } else {
401                Tx tx = getTx(message.getTransactionId());
402                tx.add(new AddMessageCommand(context) {
403                    @Override
404                    public Message getMessage() {
405                        return message;
406                    }
407                    @Override
408                    public Future<Object> run(ConnectionContext ctx) throws IOException {
409                        return destination.asyncAddQueueMessage(ctx, message);
410                    }
411
412                });
413                return AbstractMessageStore.FUTURE;
414            }
415        } else {
416            return destination.asyncAddQueueMessage(context, message);
417        }
418    }
419
420    ListenableFuture<Object> asyncAddTopicMessage(ConnectionContext context, final MessageStore destination, final Message message)
421            throws IOException {
422
423        if (message.getTransactionId() != null) {
424            if (message.getTransactionId().isXATransaction() || theStore.isConcurrentStoreAndDispatchTransactions()==false) {
425                destination.addMessage(context, message);
426                return AbstractMessageStore.FUTURE;
427            } else {
428                Tx tx = getTx(message.getTransactionId());
429                tx.add(new AddMessageCommand(context) {
430                    @Override
431                    public Message getMessage() {
432                        return message;
433                    }
434                    @Override
435                    public Future<Object> run(ConnectionContext ctx) throws IOException {
436                        return destination.asyncAddTopicMessage(ctx, message);
437                    }
438
439                });
440                return AbstractMessageStore.FUTURE;
441            }
442        } else {
443            return destination.asyncAddTopicMessage(context, message);
444        }
445    }
446
447    /**
448     * @param ack
449     * @throws IOException
450     */
451    final void removeMessage(ConnectionContext context, final MessageStore destination, final MessageAck ack)
452            throws IOException {
453
454        if (ack.isInTransaction()) {
455            if (ack.getTransactionId().isXATransaction() || theStore.isConcurrentStoreAndDispatchTransactions()== false) {
456                destination.removeMessage(context, ack);
457            } else {
458                Tx tx = getTx(ack.getTransactionId());
459                tx.add(new RemoveMessageCommand(context) {
460                    @Override
461                    public MessageAck getMessageAck() {
462                        return ack;
463                    }
464
465                    @Override
466                    public Future<Object> run(ConnectionContext ctx) throws IOException {
467                        destination.removeMessage(ctx, ack);
468                        return AbstractMessageStore.FUTURE;
469                    }
470                });
471            }
472        } else {
473            destination.removeMessage(context, ack);
474        }
475    }
476
477    final void removeAsyncMessage(ConnectionContext context, final MessageStore destination, final MessageAck ack)
478            throws IOException {
479
480        if (ack.isInTransaction()) {
481            if (ack.getTransactionId().isXATransaction() || theStore.isConcurrentStoreAndDispatchTransactions()==false) {
482                destination.removeAsyncMessage(context, ack);
483            } else {
484                Tx tx = getTx(ack.getTransactionId());
485                tx.add(new RemoveMessageCommand(context) {
486                    @Override
487                    public MessageAck getMessageAck() {
488                        return ack;
489                    }
490
491                    @Override
492                    public Future<Object> run(ConnectionContext ctx) throws IOException {
493                        destination.removeMessage(ctx, ack);
494                        return AbstractMessageStore.FUTURE;
495                    }
496                });
497            }
498        } else {
499            destination.removeAsyncMessage(context, ack);
500        }
501    }
502
503    final void acknowledge(ConnectionContext context, final TopicMessageStore destination, final String clientId, final String subscriptionName,
504                           final MessageId messageId, final MessageAck ack) throws IOException {
505
506        if (ack.isInTransaction()) {
507            if (ack.getTransactionId().isXATransaction() || theStore.isConcurrentStoreAndDispatchTransactions()== false) {
508                destination.acknowledge(context, clientId, subscriptionName, messageId, ack);
509            } else {
510                Tx tx = getTx(ack.getTransactionId());
511                tx.add(new RemoveMessageCommand(context) {
512                    public MessageAck getMessageAck() {
513                        return ack;
514                    }
515
516                    public Future<Object> run(ConnectionContext ctx) throws IOException {
517                        destination.acknowledge(ctx, clientId, subscriptionName, messageId, ack);
518                        return AbstractMessageStore.FUTURE;
519                    }
520                });
521            }
522        } else {
523            destination.acknowledge(context, clientId, subscriptionName, messageId, ack);
524        }
525    }
526
527
528    private KahaTransactionInfo getTransactionInfo(TransactionId txid) {
529        return TransactionIdConversion.convert(theStore.getTransactionIdTransformer().transform(txid));
530    }
531}