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