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.File;
020import java.io.IOException;
021import java.util.Date;
022import java.util.HashSet;
023import java.util.Set;
024import java.util.TreeSet;
025import java.util.concurrent.ConcurrentHashMap;
026
027import org.apache.activemq.broker.Broker;
028import org.apache.activemq.broker.ConnectionContext;
029import org.apache.activemq.command.Message;
030import org.apache.activemq.command.MessageAck;
031import org.apache.activemq.command.MessageId;
032import org.apache.activemq.command.TransactionId;
033import org.apache.activemq.command.XATransactionId;
034import org.apache.activemq.store.AbstractMessageStore;
035import org.apache.activemq.store.ListenableFuture;
036import org.apache.activemq.store.MessageStore;
037import org.apache.activemq.store.PersistenceAdapter;
038import org.apache.activemq.store.ProxyMessageStore;
039import org.apache.activemq.store.ProxyTopicMessageStore;
040import org.apache.activemq.store.TopicMessageStore;
041import org.apache.activemq.store.TransactionRecoveryListener;
042import org.apache.activemq.store.TransactionStore;
043import org.apache.activemq.store.kahadb.data.KahaCommitCommand;
044import org.apache.activemq.store.kahadb.data.KahaEntryType;
045import org.apache.activemq.store.kahadb.data.KahaPrepareCommand;
046import org.apache.activemq.store.kahadb.data.KahaTraceCommand;
047import org.apache.activemq.store.kahadb.disk.journal.Journal;
048import org.apache.activemq.store.kahadb.disk.journal.Location;
049import org.apache.activemq.util.DataByteArrayInputStream;
050import org.apache.activemq.util.DataByteArrayOutputStream;
051import org.apache.activemq.util.IOHelper;
052import org.slf4j.Logger;
053import org.slf4j.LoggerFactory;
054
055public class MultiKahaDBTransactionStore implements TransactionStore {
056    static final Logger LOG = LoggerFactory.getLogger(MultiKahaDBTransactionStore.class);
057    final MultiKahaDBPersistenceAdapter multiKahaDBPersistenceAdapter;
058    final ConcurrentHashMap<TransactionId, Tx> inflightTransactions = new ConcurrentHashMap<TransactionId, Tx>();
059    final Set<TransactionId> recoveredPendingCommit = new HashSet<TransactionId>();
060    private Journal journal;
061    private int journalMaxFileLength = Journal.DEFAULT_MAX_FILE_LENGTH;
062    private int journalWriteBatchSize = Journal.DEFAULT_MAX_WRITE_BATCH_SIZE;
063
064    public MultiKahaDBTransactionStore(MultiKahaDBPersistenceAdapter multiKahaDBPersistenceAdapter) {
065        this.multiKahaDBPersistenceAdapter = multiKahaDBPersistenceAdapter;
066    }
067
068    public MessageStore proxy(final TransactionStore transactionStore, MessageStore messageStore) {
069        return new ProxyMessageStore(messageStore) {
070            @Override
071            public void addMessage(ConnectionContext context, final Message send) throws IOException {
072                MultiKahaDBTransactionStore.this.addMessage(transactionStore, context, getDelegate(), send);
073            }
074
075            @Override
076            public void addMessage(ConnectionContext context, final Message send, boolean canOptimizeHint) throws IOException {
077                MultiKahaDBTransactionStore.this.addMessage(transactionStore, context, getDelegate(), send);
078            }
079
080            @Override
081            public ListenableFuture<Object> asyncAddQueueMessage(ConnectionContext context, Message message) throws IOException {
082                return MultiKahaDBTransactionStore.this.asyncAddQueueMessage(transactionStore, context, getDelegate(), message);
083            }
084
085            @Override
086            public ListenableFuture<Object> asyncAddQueueMessage(ConnectionContext context, Message message, boolean canOptimizeHint) throws IOException {
087                return MultiKahaDBTransactionStore.this.asyncAddQueueMessage(transactionStore, context, getDelegate(), message);
088            }
089
090            @Override
091            public void removeMessage(ConnectionContext context, final MessageAck ack) throws IOException {
092                MultiKahaDBTransactionStore.this.removeMessage(transactionStore, context, getDelegate(), ack);
093            }
094
095            @Override
096            public void removeAsyncMessage(ConnectionContext context, MessageAck ack) throws IOException {
097                MultiKahaDBTransactionStore.this.removeAsyncMessage(transactionStore, context, getDelegate(), ack);
098            }
099        };
100    }
101
102    public TopicMessageStore proxy(final TransactionStore transactionStore, final TopicMessageStore messageStore) {
103        return new ProxyTopicMessageStore(messageStore) {
104            @Override
105            public void addMessage(ConnectionContext context, final Message send, boolean canOptimizeHint) throws IOException {
106                MultiKahaDBTransactionStore.this.addMessage(transactionStore, context, getDelegate(), send);
107            }
108
109            @Override
110            public void addMessage(ConnectionContext context, final Message send) throws IOException {
111                MultiKahaDBTransactionStore.this.addMessage(transactionStore, context, getDelegate(), send);
112            }
113
114            @Override
115            public ListenableFuture<Object> asyncAddTopicMessage(ConnectionContext context, Message message, boolean canOptimizeHint) throws IOException {
116                return MultiKahaDBTransactionStore.this.asyncAddTopicMessage(transactionStore, context, getDelegate(), message);
117            }
118
119            @Override
120            public ListenableFuture<Object> asyncAddTopicMessage(ConnectionContext context, Message message) throws IOException {
121                return MultiKahaDBTransactionStore.this.asyncAddTopicMessage(transactionStore, context, getDelegate(), message);
122            }
123
124            @Override
125            public void removeMessage(ConnectionContext context, final MessageAck ack) throws IOException {
126                MultiKahaDBTransactionStore.this.removeMessage(transactionStore, context, getDelegate(), ack);
127            }
128
129            @Override
130            public void removeAsyncMessage(ConnectionContext context, MessageAck ack) throws IOException {
131                MultiKahaDBTransactionStore.this.removeAsyncMessage(transactionStore, context, getDelegate(), ack);
132            }
133
134            @Override
135            public void acknowledge(ConnectionContext context, String clientId, String subscriptionName,
136                                    MessageId messageId, MessageAck ack) throws IOException {
137                MultiKahaDBTransactionStore.this.acknowledge(transactionStore, context, (TopicMessageStore) getDelegate(), clientId,
138                        subscriptionName, messageId, ack);
139            }
140        };
141    }
142
143    public void deleteAllMessages() {
144        IOHelper.deleteChildren(getDirectory());
145    }
146
147    public int getJournalMaxFileLength() {
148        return journalMaxFileLength;
149    }
150
151    public void setJournalMaxFileLength(int journalMaxFileLength) {
152        this.journalMaxFileLength = journalMaxFileLength;
153    }
154
155    public int getJournalMaxWriteBatchSize() {
156        return journalWriteBatchSize;
157    }
158
159    public void setJournalMaxWriteBatchSize(int journalWriteBatchSize) {
160        this.journalWriteBatchSize = journalWriteBatchSize;
161    }
162
163    public class Tx {
164        private final Set<TransactionStore> stores = new HashSet<TransactionStore>();
165        private int prepareLocationId = 0;
166
167        public void trackStore(TransactionStore store) {
168            stores.add(store);
169        }
170
171        public Set<TransactionStore> getStores() {
172            return stores;
173        }
174
175        public void trackPrepareLocation(Location location) {
176            this.prepareLocationId = location.getDataFileId();
177        }
178
179        public int getPreparedLocationId() {
180            return prepareLocationId;
181        }
182    }
183
184    public Tx getTx(TransactionId txid) {
185        Tx tx = inflightTransactions.get(txid);
186        if (tx == null) {
187            tx = new Tx();
188            inflightTransactions.put(txid, tx);
189        }
190        return tx;
191    }
192
193    public Tx removeTx(TransactionId txid) {
194        return inflightTransactions.remove(txid);
195    }
196
197    @Override
198    public void prepare(TransactionId txid) throws IOException {
199        Tx tx = getTx(txid);
200        for (TransactionStore store : tx.getStores()) {
201            store.prepare(txid);
202        }
203    }
204
205    @Override
206    public void commit(TransactionId txid, boolean wasPrepared, Runnable preCommit, Runnable postCommit)
207            throws IOException {
208
209        if (preCommit != null) {
210            preCommit.run();
211        }
212
213        Tx tx = getTx(txid);
214        if (wasPrepared) {
215            for (TransactionStore store : tx.getStores()) {
216                store.commit(txid, true, null, null);
217            }
218        } else {
219            // can only do 1pc on a single store
220            if (tx.getStores().size() == 1) {
221                for (TransactionStore store : tx.getStores()) {
222                    store.commit(txid, false, null, null);
223                }
224            } else {
225                // need to do local 2pc
226                for (TransactionStore store : tx.getStores()) {
227                    store.prepare(txid);
228                }
229                persistOutcome(tx, txid);
230                for (TransactionStore store : tx.getStores()) {
231                    store.commit(txid, true, null, null);
232                }
233                persistCompletion(txid);
234            }
235        }
236        removeTx(txid);
237        if (postCommit != null) {
238            postCommit.run();
239        }
240    }
241
242    public void persistOutcome(Tx tx, TransactionId txid) throws IOException {
243        tx.trackPrepareLocation(store(new KahaPrepareCommand().setTransactionInfo(TransactionIdConversion.convert(multiKahaDBPersistenceAdapter.transactionIdTransformer.transform(txid)))));
244    }
245
246    public void persistCompletion(TransactionId txid) throws IOException {
247        store(new KahaCommitCommand().setTransactionInfo(TransactionIdConversion.convert(multiKahaDBPersistenceAdapter.transactionIdTransformer.transform(txid))));
248    }
249
250    private Location store(JournalCommand<?> data) throws IOException {
251        int size = data.serializedSizeFramed();
252        DataByteArrayOutputStream os = new DataByteArrayOutputStream(size + 1);
253        os.writeByte(data.type().getNumber());
254        data.writeFramed(os);
255        Location location = journal.write(os.toByteSequence(), true);
256        journal.setLastAppendLocation(location);
257        return location;
258    }
259
260    @Override
261    public void rollback(TransactionId txid) throws IOException {
262        Tx tx = removeTx(txid);
263        if (tx != null) {
264            for (TransactionStore store : tx.getStores()) {
265                store.rollback(txid);
266            }
267        }
268    }
269
270    @Override
271    public void start() throws Exception {
272        journal = new Journal() {
273            @Override
274            protected void cleanup() {
275                super.cleanup();
276                txStoreCleanup();
277            }
278        };
279        journal.setDirectory(getDirectory());
280        journal.setMaxFileLength(journalMaxFileLength);
281        journal.setWriteBatchSize(journalWriteBatchSize);
282        IOHelper.mkdirs(journal.getDirectory());
283        journal.start();
284        recoverPendingLocalTransactions();
285        store(new KahaTraceCommand().setMessage("LOADED " + new Date()));
286    }
287
288    private void txStoreCleanup() {
289        Set<Integer> knownDataFileIds = new TreeSet<Integer>(journal.getFileMap().keySet());
290        for (Tx tx : inflightTransactions.values()) {
291            knownDataFileIds.remove(tx.getPreparedLocationId());
292        }
293        try {
294            journal.removeDataFiles(knownDataFileIds);
295        } catch (Exception e) {
296            LOG.error(this + ", Failed to remove tx journal datafiles " + knownDataFileIds);
297        }
298    }
299
300    private File getDirectory() {
301        return new File(multiKahaDBPersistenceAdapter.getDirectory(), "txStore");
302    }
303
304    @Override
305    public void stop() throws Exception {
306        journal.close();
307        journal = null;
308    }
309
310    private void recoverPendingLocalTransactions() throws IOException {
311        Location location = journal.getNextLocation(null);
312        while (location != null) {
313            process(load(location));
314            location = journal.getNextLocation(location);
315        }
316        recoveredPendingCommit.addAll(inflightTransactions.keySet());
317        LOG.info("pending local transactions: " + recoveredPendingCommit);
318    }
319
320    public JournalCommand<?> load(Location location) throws IOException {
321        DataByteArrayInputStream is = new DataByteArrayInputStream(journal.read(location));
322        byte readByte = is.readByte();
323        KahaEntryType type = KahaEntryType.valueOf(readByte);
324        if (type == null) {
325            throw new IOException("Could not load journal record. Invalid location: " + location);
326        }
327        JournalCommand<?> message = (JournalCommand<?>) type.createMessage();
328        message.mergeFramed(is);
329        return message;
330    }
331
332    public void process(JournalCommand<?> command) throws IOException {
333        switch (command.type()) {
334            case KAHA_PREPARE_COMMAND:
335                KahaPrepareCommand prepareCommand = (KahaPrepareCommand) command;
336                getTx(TransactionIdConversion.convert(prepareCommand.getTransactionInfo()));
337                break;
338            case KAHA_COMMIT_COMMAND:
339                KahaCommitCommand commitCommand = (KahaCommitCommand) command;
340                removeTx(TransactionIdConversion.convert(commitCommand.getTransactionInfo()));
341                break;
342            case KAHA_TRACE_COMMAND:
343                break;
344            default:
345                throw new IOException("Unexpected command in transaction journal: " + command);
346        }
347    }
348
349
350    @Override
351    public synchronized void recover(final TransactionRecoveryListener listener) throws IOException {
352
353        for (final PersistenceAdapter adapter : multiKahaDBPersistenceAdapter.adapters) {
354            adapter.createTransactionStore().recover(new TransactionRecoveryListener() {
355                @Override
356                public void recover(XATransactionId xid, Message[] addedMessages, MessageAck[] acks) {
357                    try {
358                        getTx(xid).trackStore(adapter.createTransactionStore());
359                    } catch (IOException e) {
360                        LOG.error("Failed to access transaction store: " + adapter + " for prepared xa tid: " + xid, e);
361                    }
362                    listener.recover(xid, addedMessages, acks);
363                }
364            });
365        }
366
367        try {
368            Broker broker = multiKahaDBPersistenceAdapter.getBrokerService().getBroker();
369            // force completion of local xa
370            for (TransactionId txid : broker.getPreparedTransactions(null)) {
371                if (multiKahaDBPersistenceAdapter.isLocalXid(txid)) {
372                    try {
373                        if (recoveredPendingCommit.contains(txid)) {
374                            LOG.info("delivering pending commit outcome for tid: " + txid);
375                            broker.commitTransaction(null, txid, false);
376
377                        } else {
378                            LOG.info("delivering rollback outcome to store for tid: " + txid);
379                            broker.forgetTransaction(null, txid);
380                        }
381                        persistCompletion(txid);
382                    } catch (Exception ex) {
383                        LOG.error("failed to deliver pending outcome for tid: " + txid, ex);
384                    }
385                }
386            }
387        } catch (Exception e) {
388            LOG.error("failed to resolve pending local transactions", e);
389        }
390    }
391
392    void addMessage(final TransactionStore transactionStore, ConnectionContext context, final MessageStore destination, final Message message)
393            throws IOException {
394        if (message.getTransactionId() != null) {
395            getTx(message.getTransactionId()).trackStore(transactionStore);
396        }
397        destination.addMessage(context, message);
398    }
399
400    ListenableFuture<Object> asyncAddQueueMessage(final TransactionStore transactionStore, ConnectionContext context, final MessageStore destination, final Message message)
401            throws IOException {
402        if (message.getTransactionId() != null) {
403            getTx(message.getTransactionId()).trackStore(transactionStore);
404            destination.addMessage(context, message);
405            return AbstractMessageStore.FUTURE;
406        } else {
407            return destination.asyncAddQueueMessage(context, message);
408        }
409    }
410
411    ListenableFuture<Object> asyncAddTopicMessage(final TransactionStore transactionStore, ConnectionContext context, final MessageStore destination, final Message message)
412            throws IOException {
413
414        if (message.getTransactionId() != null) {
415            getTx(message.getTransactionId()).trackStore(transactionStore);
416            destination.addMessage(context, message);
417            return AbstractMessageStore.FUTURE;
418        } else {
419            return destination.asyncAddTopicMessage(context, message);
420        }
421    }
422
423    final void removeMessage(final TransactionStore transactionStore, ConnectionContext context, final MessageStore destination, final MessageAck ack)
424            throws IOException {
425        if (ack.getTransactionId() != null) {
426            getTx(ack.getTransactionId()).trackStore(transactionStore);
427        }
428        destination.removeMessage(context, ack);
429    }
430
431    final void removeAsyncMessage(final TransactionStore transactionStore, ConnectionContext context, final MessageStore destination, final MessageAck ack)
432            throws IOException {
433        if (ack.getTransactionId() != null) {
434            getTx(ack.getTransactionId()).trackStore(transactionStore);
435        }
436        destination.removeAsyncMessage(context, ack);
437    }
438
439    final void acknowledge(final TransactionStore transactionStore, ConnectionContext context, final TopicMessageStore destination,
440                           final String clientId, final String subscriptionName,
441                           final MessageId messageId, final MessageAck ack) throws IOException {
442        if (ack.getTransactionId() != null) {
443            getTx(ack.getTransactionId()).trackStore(transactionStore);
444        }
445        destination.acknowledge(context, clientId, subscriptionName, messageId, ack);
446    }
447
448}