activemq-kahadb-store/src/main/java/org/apache/activemq/store/kahadb/KahaDBTransactionStore.java

/**
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.activemq.store.kahadb;

import java.io.DataInputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.concurrent.Future;

import org.apache.activemq.broker.ConnectionContext;
import org.apache.activemq.command.ConsumerId;
import org.apache.activemq.command.Message;
import org.apache.activemq.command.MessageAck;
import org.apache.activemq.command.MessageId;
import org.apache.activemq.command.TransactionId;
import org.apache.activemq.command.XATransactionId;
import org.apache.activemq.protobuf.Buffer;
import org.apache.activemq.store.AbstractMessageStore;
import org.apache.activemq.store.ListenableFuture;
import org.apache.activemq.store.MessageStore;
import org.apache.activemq.store.ProxyMessageStore;
import org.apache.activemq.store.ProxyTopicMessageStore;
import org.apache.activemq.store.TopicMessageStore;
import org.apache.activemq.store.TransactionRecoveryListener;
import org.apache.activemq.store.TransactionStore;
import org.apache.activemq.store.kahadb.MessageDatabase.Operation;
import org.apache.activemq.store.kahadb.data.KahaCommitCommand;
import org.apache.activemq.store.kahadb.data.KahaPrepareCommand;
import org.apache.activemq.store.kahadb.data.KahaRollbackCommand;
import org.apache.activemq.store.kahadb.data.KahaTransactionInfo;
import org.apache.activemq.wireformat.WireFormat;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * Provides a TransactionStore implementation that can create transaction aware
 * MessageStore objects from non transaction aware MessageStore objects.
 *
 *
 */
public class KahaDBTransactionStore implements TransactionStore {
    static final Logger LOG = LoggerFactory.getLogger(KahaDBTransactionStore.class);
    private final KahaDBStore theStore;

    public KahaDBTransactionStore(KahaDBStore theStore) {
        this.theStore = theStore;
    }

    private WireFormat wireFormat(){
      return this.theStore.wireFormat;
    }

    public class Tx {
        private final List<AddMessageCommand> messages = Collections.synchronizedList(new ArrayList<>());

        private final List<RemoveMessageCommand> acks = Collections.synchronizedList(new ArrayList<>());

        public void add(AddMessageCommand msg) {
            messages.add(msg);
        }

        public void add(RemoveMessageCommand ack) {
            acks.add(ack);
        }

        public Message[] getMessages() {
            Message rc[] = new Message[messages.size()];
            int count = 0;
            for (AddMessageCommand cmd : messages) {
                rc[count++] = cmd.getMessage();
            }
            return rc;
        }

        public MessageAck[] getAcks() {
            MessageAck rc[] = new MessageAck[acks.size()];
            int count = 0;
            for (RemoveMessageCommand cmd : acks) {
                rc[count++] = cmd.getMessageAck();
            }
            return rc;
        }

        /**
         * @return true if something to commit
         * @throws IOException
         */
        public List<Future<Object>> commit() throws IOException {
            List<Future<Object>> results = new ArrayList<>();
            // Do all the message adds.
            for (AddMessageCommand cmd : messages) {
                results.add(cmd.run());

            }
            // And removes..
            for (RemoveMessageCommand cmd : acks) {
                cmd.run();
                results.add(cmd.run());
            }

            return results;
        }
    }

    public abstract static class AddMessageCommand {
        private final ConnectionContext ctx;
        AddMessageCommand(ConnectionContext ctx) {
            this.ctx = ctx;
        }
        abstract Message getMessage();
        Future<Object> run() throws IOException {
            return run(this.ctx);
        }
        abstract Future<Object> run(ConnectionContext ctx) throws IOException;
    }

    public abstract static class RemoveMessageCommand {

        private final ConnectionContext ctx;
        RemoveMessageCommand(ConnectionContext ctx) {
            this.ctx = ctx;
        }
        abstract MessageAck getMessageAck();
        Future<Object> run() throws IOException {
            return run(this.ctx);
        }
        abstract Future<Object> run(ConnectionContext context) throws IOException;
    }

    public MessageStore proxy(MessageStore messageStore) {
        return new ProxyMessageStore(messageStore) {
            @Override
            public void addMessage(ConnectionContext context, final Message send) throws IOException {
                KahaDBTransactionStore.this.addMessage(context, getDelegate(), send);
            }

            @Override
            public void addMessage(ConnectionContext context, final Message send, boolean canOptimize) throws IOException {
                KahaDBTransactionStore.this.addMessage(context, getDelegate(), send);
            }

            @Override
            public ListenableFuture<Object> asyncAddQueueMessage(ConnectionContext context, Message message) throws IOException {
                return KahaDBTransactionStore.this.asyncAddQueueMessage(context, getDelegate(), message);
            }

            @Override
            public ListenableFuture<Object> asyncAddQueueMessage(ConnectionContext context, Message message, boolean canOptimize) throws IOException {
                return KahaDBTransactionStore.this.asyncAddQueueMessage(context, getDelegate(), message);
            }

            @Override
            public void removeMessage(ConnectionContext context, final MessageAck ack) throws IOException {
                KahaDBTransactionStore.this.removeMessage(context, getDelegate(), ack);
            }

            @Override
            public void removeAsyncMessage(ConnectionContext context, MessageAck ack) throws IOException {
                KahaDBTransactionStore.this.removeAsyncMessage(context, getDelegate(), ack);
            }
        };
    }

    public TopicMessageStore proxy(TopicMessageStore messageStore) {
        return new ProxyTopicMessageStore(messageStore) {
            @Override
            public void addMessage(ConnectionContext context, final Message send) throws IOException {
                KahaDBTransactionStore.this.addMessage(context, getDelegate(), send);
            }

            @Override
            public void addMessage(ConnectionContext context, final Message send, boolean canOptimize) throws IOException {
                KahaDBTransactionStore.this.addMessage(context, getDelegate(), send);
            }

            @Override
            public ListenableFuture<Object> asyncAddTopicMessage(ConnectionContext context, Message message) throws IOException {
                return KahaDBTransactionStore.this.asyncAddTopicMessage(context, getDelegate(), message);
            }

            @Override
            public ListenableFuture<Object> asyncAddTopicMessage(ConnectionContext context, Message message, boolean canOptimize) throws IOException {
                return KahaDBTransactionStore.this.asyncAddTopicMessage(context, getDelegate(), message);
            }

            @Override
            public void removeMessage(ConnectionContext context, final MessageAck ack) throws IOException {
                KahaDBTransactionStore.this.removeMessage(context, getDelegate(), ack);
            }

            @Override
            public void removeAsyncMessage(ConnectionContext context, MessageAck ack) throws IOException {
                KahaDBTransactionStore.this.removeAsyncMessage(context, getDelegate(), ack);
            }

            @Override
            public void acknowledge(ConnectionContext context, String clientId, String subscriptionName,
                            MessageId messageId, MessageAck ack) throws IOException {
                KahaDBTransactionStore.this.acknowledge(context, (TopicMessageStore)getDelegate(), clientId,
                        subscriptionName, messageId, ack);
            }

        };
    }

    /**
     * @throws IOException
     * @see org.apache.activemq.store.TransactionStore#prepare(TransactionId)
     */
    @Override
    public void prepare(TransactionId txid) throws IOException {
        KahaTransactionInfo info = getTransactionInfo(txid);
        theStore.store(new KahaPrepareCommand().setTransactionInfo(info), true, null, null);
    }

    @Override
    public void commit(TransactionId txid, boolean wasPrepared, final Runnable preCommit, Runnable postCommit)
            throws IOException {
        if (txid != null) {
            KahaTransactionInfo info = getTransactionInfo(txid);
            if (preCommit != null) {
                preCommit.run();
            }
            theStore.store(new KahaCommitCommand().setTransactionInfo(info), theStore.isEnableJournalDiskSyncs(), null, postCommit);
            forgetRecoveredAcks(txid, false);
        } else {
           LOG.error("Null transaction passed on commit");
        }
    }

    /**
     * @throws IOException
     * @see org.apache.activemq.store.TransactionStore#rollback(TransactionId)
     */
    @Override
    public void rollback(TransactionId txid) throws IOException {
        KahaTransactionInfo info = getTransactionInfo(txid);
        theStore.store(new KahaRollbackCommand().setTransactionInfo(info), theStore.isEnableJournalDiskSyncs(), null, null);
        forgetRecoveredAcks(txid, true);
    }

    protected void forgetRecoveredAcks(TransactionId txid, boolean isRollback) throws IOException {
        if (txid.isXATransaction()) {
            XATransactionId xaTid = ((XATransactionId) txid);
            theStore.forgetRecoveredAcks(xaTid.getPreparedAcks(), isRollback);
        }
    }

    @Override
    public void start() throws Exception {
    }

    @Override
    public void stop() throws Exception {
    }

    @Override
    public synchronized void recover(TransactionRecoveryListener listener) throws IOException {
        for (Map.Entry<TransactionId, List<Operation>> entry : theStore.preparedTransactions.entrySet()) {
            XATransactionId xid = (XATransactionId) entry.getKey();
            ArrayList<Message> messageList = new ArrayList<>();
            ArrayList<MessageAck> ackList = new ArrayList<>();

            for (Operation op : entry.getValue()) {
                if (op.getClass() == MessageDatabase.AddOperation.class) {
                    MessageDatabase.AddOperation addOp = (MessageDatabase.AddOperation) op;
                    Message msg = (Message) wireFormat().unmarshal(new DataInputStream(addOp.getCommand().getMessage()
                            .newInput()));
                    messageList.add(msg);
                } else {
                    MessageDatabase.RemoveOperation rmOp = (MessageDatabase.RemoveOperation) op;
                    Buffer ackb = rmOp.getCommand().getAck();
                    MessageAck ack = (MessageAck) wireFormat().unmarshal(new DataInputStream(ackb.newInput()));
                    // allow the ack to be tracked back to its durable sub
                    ConsumerId subKey = new ConsumerId();
                    subKey.setConnectionId(rmOp.getCommand().getSubscriptionKey());
                    ack.setConsumerId(subKey);
                    ackList.add(ack);
                }
            }

            Message[] addedMessages = new Message[messageList.size()];
            MessageAck[] acks = new MessageAck[ackList.size()];
            messageList.toArray(addedMessages);
            ackList.toArray(acks);
            xid.setPreparedAcks(ackList);
            theStore.trackRecoveredAcks(ackList);
            listener.recover(xid, addedMessages, acks);
        }
    }

    /**
     * @param message
     * @throws IOException
     */
    void addMessage(ConnectionContext context, final MessageStore destination, final Message message)
            throws IOException {
        destination.addMessage(context, message);
    }

    ListenableFuture<Object> asyncAddQueueMessage(ConnectionContext context, final MessageStore destination, final Message message)
            throws IOException {
        if (message.getTransactionId() != null) {
            destination.addMessage(context, message);
            return AbstractMessageStore.FUTURE;
        } else {
            return destination.asyncAddQueueMessage(context, message);
        }
    }

    ListenableFuture<Object> asyncAddTopicMessage(ConnectionContext context, final MessageStore destination, final Message message)
            throws IOException {

        if (message.getTransactionId() != null) {
            destination.addMessage(context, message);
            return AbstractMessageStore.FUTURE;
        } else {
            return destination.asyncAddTopicMessage(context, message);
        }
    }

    /**
     * @param ack
     * @throws IOException
     */
    final void removeMessage(ConnectionContext context, final MessageStore destination, final MessageAck ack)
            throws IOException {
        destination.removeMessage(context, ack);
    }

    final void removeAsyncMessage(ConnectionContext context, final MessageStore destination, final MessageAck ack)
            throws IOException {
        destination.removeAsyncMessage(context, ack);
    }

    final void acknowledge(ConnectionContext context, final TopicMessageStore destination, final String clientId, final String subscriptionName,
                           final MessageId messageId, final MessageAck ack) throws IOException {
        destination.acknowledge(context, clientId, subscriptionName, messageId, ack);
    }


    private KahaTransactionInfo getTransactionInfo(TransactionId txid) {
        return TransactionIdConversion.convert(theStore.getTransactionIdTransformer().transform(txid));
    }
}