content: /**
   * 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.File;
  import java.io.IOException;
  import java.util.Date;
  import java.util.Map;
  import java.util.Set;
  import java.util.TreeSet;
  import java.util.concurrent.ConcurrentHashMap;
  import java.util.concurrent.ConcurrentMap;
  import java.util.concurrent.atomic.AtomicBoolean;
  
  import org.apache.activemq.broker.Broker;
  import org.apache.activemq.broker.ConnectionContext;
  import org.apache.activemq.broker.region.BaseDestination;
  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.store.AbstractMessageStore;
  import org.apache.activemq.store.IndexListener;
  import org.apache.activemq.store.ListenableFuture;
  import org.apache.activemq.store.MessageStore;
  import org.apache.activemq.store.PersistenceAdapter;
  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.data.KahaCommitCommand;
  import org.apache.activemq.store.kahadb.data.KahaEntryType;
  import org.apache.activemq.store.kahadb.data.KahaPrepareCommand;
  import org.apache.activemq.store.kahadb.data.KahaTraceCommand;
  import org.apache.activemq.store.kahadb.disk.journal.DataFile;
  import org.apache.activemq.store.kahadb.disk.journal.Journal;
  import org.apache.activemq.store.kahadb.disk.journal.Location;
  import org.apache.activemq.usage.StoreUsage;
  import org.apache.activemq.util.DataByteArrayInputStream;
  import org.apache.activemq.util.DataByteArrayOutputStream;
  import org.apache.activemq.util.IOHelper;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  public class MultiKahaDBTransactionStore implements TransactionStore {
      static final Logger LOG = LoggerFactory.getLogger(MultiKahaDBTransactionStore.class);
      final MultiKahaDBPersistenceAdapter multiKahaDBPersistenceAdapter;
      final ConcurrentMap<TransactionId, Tx> inflightTransactions = new ConcurrentHashMap<TransactionId, Tx>();
      final ConcurrentMap<TransactionId, Tx> pendingCommit = new ConcurrentHashMap<TransactionId, Tx>();
      private Journal journal;
      private int journalMaxFileLength = Journal.DEFAULT_MAX_FILE_LENGTH;
      private int journalWriteBatchSize = Journal.DEFAULT_MAX_WRITE_BATCH_SIZE;
      private final AtomicBoolean started = new AtomicBoolean(false);
      private final AtomicBoolean recovered = new AtomicBoolean(false);
      private long journalCleanupInterval = Journal.DEFAULT_CLEANUP_INTERVAL;
      private boolean checkForCorruption = true;
      private AtomicBoolean corruptJournalDetected = new AtomicBoolean(false);
  
      public MultiKahaDBTransactionStore(MultiKahaDBPersistenceAdapter multiKahaDBPersistenceAdapter) {
          this.multiKahaDBPersistenceAdapter = multiKahaDBPersistenceAdapter;
      }
  
      public MessageStore proxy(final TransactionStore transactionStore, MessageStore messageStore) {
          return new ProxyMessageStore(messageStore) {
              @Override
              public void addMessage(ConnectionContext context, final Message send) throws IOException {
                  MultiKahaDBTransactionStore.this.addMessage(transactionStore, context, getDelegate(), send);
              }
  
              @Override
              public void addMessage(ConnectionContext context, final Message send, boolean canOptimizeHint) throws IOException {
                  MultiKahaDBTransactionStore.this.addMessage(transactionStore, context, getDelegate(), send);
              }
  
              @Override
              public ListenableFuture<Object> asyncAddQueueMessage(ConnectionContext context, Message message) throws IOException {
                  return MultiKahaDBTransactionStore.this.asyncAddQueueMessage(transactionStore, context, getDelegate(), message);
              }
  
              @Override
              public ListenableFuture<Object> asyncAddQueueMessage(ConnectionContext context, Message message, boolean canOptimizeHint) throws IOException {
                  return MultiKahaDBTransactionStore.this.asyncAddQueueMessage(transactionStore, context, getDelegate(), message);
              }
  
              @Override
              public void removeMessage(ConnectionContext context, final MessageAck ack) throws IOException {
                  MultiKahaDBTransactionStore.this.removeMessage(transactionStore, context, getDelegate(), ack);
              }
  
              @Override
              public void removeAsyncMessage(ConnectionContext context, MessageAck ack) throws IOException {
                  MultiKahaDBTransactionStore.this.removeAsyncMessage(transactionStore, context, getDelegate(), ack);
              }
  
              @Override
              public void registerIndexListener(IndexListener indexListener) {
                  getDelegate().registerIndexListener(indexListener);
                  try {
                      if (indexListener instanceof BaseDestination) {
                          // update queue storeUsage
                          Object matchingPersistenceAdapter = multiKahaDBPersistenceAdapter.destinationMap.chooseValue(getDelegate().getDestination());
                          if (matchingPersistenceAdapter instanceof FilteredKahaDBPersistenceAdapter) {
                              FilteredKahaDBPersistenceAdapter filteredAdapter = (FilteredKahaDBPersistenceAdapter) matchingPersistenceAdapter;
                              if (filteredAdapter.getUsage() != null && filteredAdapter.getPersistenceAdapter() instanceof KahaDBPersistenceAdapter) {
                                  StoreUsage storeUsage = filteredAdapter.getUsage();
                                  storeUsage.setStore(filteredAdapter.getPersistenceAdapter());
                                  storeUsage.setParent(multiKahaDBPersistenceAdapter.getBrokerService().getSystemUsage().getStoreUsage());
                                  ((BaseDestination) indexListener).getSystemUsage().setStoreUsage(storeUsage);
                              }
                          }
                      }
                  } catch (Exception ignored) {
                      LOG.warn("Failed to set mKahaDB destination store usage", ignored);
                  }
              }
          };
      }
  
      public TopicMessageStore proxy(final TransactionStore transactionStore, final TopicMessageStore messageStore) {
          return new ProxyTopicMessageStore(messageStore) {
              @Override
              public void addMessage(ConnectionContext context, final Message send, boolean canOptimizeHint) throws IOException {
                  MultiKahaDBTransactionStore.this.addMessage(transactionStore, context, getDelegate(), send);
              }
  
              @Override
              public void addMessage(ConnectionContext context, final Message send) throws IOException {
                  MultiKahaDBTransactionStore.this.addMessage(transactionStore, context, getDelegate(), send);
              }
  
              @Override
              public ListenableFuture<Object> asyncAddTopicMessage(ConnectionContext context, Message message, boolean canOptimizeHint) throws IOException {
                  return MultiKahaDBTransactionStore.this.asyncAddTopicMessage(transactionStore, context, getDelegate(), message);
              }
  
              @Override
              public ListenableFuture<Object> asyncAddTopicMessage(ConnectionContext context, Message message) throws IOException {
                  return MultiKahaDBTransactionStore.this.asyncAddTopicMessage(transactionStore, context, getDelegate(), message);
              }
  
              @Override
              public void removeMessage(ConnectionContext context, final MessageAck ack) throws IOException {
                  MultiKahaDBTransactionStore.this.removeMessage(transactionStore, context, getDelegate(), ack);
              }
  
              @Override
              public void removeAsyncMessage(ConnectionContext context, MessageAck ack) throws IOException {
                  MultiKahaDBTransactionStore.this.removeAsyncMessage(transactionStore, context, getDelegate(), ack);
              }
  
              @Override
              public void acknowledge(ConnectionContext context, String clientId, String subscriptionName,
                                      MessageId messageId, MessageAck ack) throws IOException {
                  MultiKahaDBTransactionStore.this.acknowledge(transactionStore, context, (TopicMessageStore) getDelegate(), clientId,
                          subscriptionName, messageId, ack);
              }
          };
      }
  
      public void deleteAllMessages() {
          IOHelper.deleteChildren(getDirectory());
      }
  
      public int getJournalMaxFileLength() {
          return journalMaxFileLength;
      }
  
      public void setJournalMaxFileLength(int journalMaxFileLength) {
          this.journalMaxFileLength = journalMaxFileLength;
      }
  
      public int getJournalMaxWriteBatchSize() {
          return journalWriteBatchSize;
      }
  
      public void setJournalMaxWriteBatchSize(int journalWriteBatchSize) {
          this.journalWriteBatchSize = journalWriteBatchSize;
      }
  
      public void setJournalCleanupInterval(long journalCleanupInterval) {
          this.journalCleanupInterval = journalCleanupInterval;
      }
  
      public long getJournalCleanupInterval() {
          return journalCleanupInterval;
      }
  
      public void setCheckForCorruption(boolean checkForCorruption) {
          this.checkForCorruption = checkForCorruption;
      }
  
      public boolean isCheckForCorruption() {
          return checkForCorruption;
      }
  
      private static final XATransactionId NULL_XA_TRANSACTION_ID = new XATransactionId();
      public class Tx {
          private final ConcurrentHashMap<TransactionStore, TransactionId> stores = new ConcurrentHashMap<TransactionStore, TransactionId>();
          private int prepareLocationId = 0;
  
          public void trackStore(TransactionStore store, XATransactionId xid) {
              stores.put(store, xid);
          }
  
          public void trackStore(TransactionStore store) {
              stores.putIfAbsent(store, NULL_XA_TRANSACTION_ID);
          }
  
          public Map<TransactionStore, TransactionId> getStoresMap() {
              return stores;
          }
  
          public Set<TransactionStore> getStores() {
              return stores.keySet();
          }
  
          public void trackPrepareLocation(Location location) {
              this.prepareLocationId = location.getDataFileId();
          }
  
          public int getPreparedLocationId() {
              return prepareLocationId;
          }
      }
  
      public Tx getTx(TransactionId txid) {
          Tx tx = inflightTransactions.get(txid);
          if (tx == null) {
              final Tx val = new Tx();
              tx = inflightTransactions.putIfAbsent(txid, val);
              if (tx == null) {
                  // we won
                  tx = val;
              }
          }
          return tx;
      }
  
      public Tx removeTx(TransactionId txid) {
          return inflightTransactions.remove(txid);
      }
  
      @Override
      public void prepare(TransactionId txid) throws IOException {
          Tx tx = getTx(txid);
          for (TransactionStore store : tx.getStores()) {
              store.prepare(txid);
          }
      }
  
      @Override
      public void commit(TransactionId txid, boolean wasPrepared, Runnable preCommit, Runnable postCommit)
              throws IOException {
  
          if (preCommit != null) {
              preCommit.run();
          }
  
          Tx tx = getTx(txid);
          if (wasPrepared) {
              for (Map.Entry<TransactionStore, TransactionId> storeTx : tx.getStoresMap().entrySet()) {
                  TransactionId recovered = storeTx.getValue();
                  if (recovered != null && recovered != NULL_XA_TRANSACTION_ID ) {
                      storeTx.getKey().commit(recovered, true, null, null);
                  } else {
                      storeTx.getKey().commit(txid, true, null, null);
                  }
              }
          } else {
              // can only do 1pc on a single store
              if (tx.getStores().size() == 1) {
                  for (TransactionStore store : tx.getStores()) {
                      store.commit(txid, false, null, null);
                  }
              } else {
                  // need to do local 2pc
                  for (TransactionStore store : tx.getStores()) {
                      store.prepare(txid);
                  }
                  persistOutcome(tx, txid);
                  for (TransactionStore store : tx.getStores()) {
                      store.commit(txid, true, null, null);
                  }
                  persistCompletion(txid);
              }
          }
          removeTx(txid);
          if (postCommit != null) {
              postCommit.run();
          }
      }
  
      public void persistOutcome(Tx tx, TransactionId txid) throws IOException {
          tx.trackPrepareLocation(store(new KahaPrepareCommand().setTransactionInfo(TransactionIdConversion.convert(multiKahaDBPersistenceAdapter.transactionIdTransformer.transform(txid)))));
          pendingCommit.put(txid, tx);
      }
  
      public void persistCompletion(TransactionId txid) throws IOException {
          store(new KahaCommitCommand().setTransactionInfo(TransactionIdConversion.convert(multiKahaDBPersistenceAdapter.transactionIdTransformer.transform(txid))));
          pendingCommit.remove(txid);
      }
  
      private Location store(JournalCommand<?> data) throws IOException {
          int size = data.serializedSizeFramed();
          DataByteArrayOutputStream os = new DataByteArrayOutputStream(size + 1);
          os.writeByte(data.type().getNumber());
          data.writeFramed(os);
          Location location = journal.write(os.toByteSequence(), true);
          journal.setLastAppendLocation(location);
          return location;
      }
  
      @Override
      public void rollback(TransactionId txid) throws IOException {
          Tx tx = removeTx(txid);
          if (tx != null) {
              for (Map.Entry<TransactionStore, TransactionId> storeTx : tx.getStoresMap().entrySet()) {
                  TransactionId recovered = storeTx.getValue();
                  if (recovered != null && recovered != NULL_XA_TRANSACTION_ID) {
                      storeTx.getKey().rollback(recovered);
                  } else {
                      storeTx.getKey().rollback(txid);
                  }
              }
          }
      }
  
      @Override
      public void start() throws Exception {
          if (started.compareAndSet(false, true)) {
              journal = new Journal() {
                  @Override
                  public void cleanup() {
                      super.cleanup();
                      txStoreCleanup();
                  }
              };
              journal.setDirectory(getDirectory());
              journal.setMaxFileLength(journalMaxFileLength);
              journal.setWriteBatchSize(journalWriteBatchSize);
              journal.setCleanupInterval(journalCleanupInterval);
              journal.setCheckForCorruptionOnStartup(checkForCorruption);
              journal.setChecksum(checkForCorruption);
              IOHelper.mkdirs(journal.getDirectory());
              journal.start();
              recoverPendingLocalTransactions();
              recovered.set(true);
              loaded();
          }
      }
  
      private void loaded() throws IOException {
          store(new KahaTraceCommand().setMessage("LOADED " + new Date()));
      }
  
      private void txStoreCleanup() {
          if (!recovered.get() || corruptJournalDetected.get()) {
              return;
          }
          Set<Integer> knownDataFileIds = new TreeSet<Integer>(journal.getFileMap().keySet());
          for (Tx tx : inflightTransactions.values()) {
              knownDataFileIds.remove(tx.getPreparedLocationId());
          }
          for (Tx tx : pendingCommit.values()) {
              knownDataFileIds.remove(tx.getPreparedLocationId());
          }
          try {
              journal.removeDataFiles(knownDataFileIds);
          } catch (Exception e) {
              LOG.error(this + ", Failed to remove tx journal datafiles " + knownDataFileIds);
          }
      }
  
      private File getDirectory() {
          return new File(multiKahaDBPersistenceAdapter.getDirectory(), "txStore");
      }
  
      @Override
      public void stop() throws Exception {
          if (started.compareAndSet(true, false) && journal != null) {
              journal.close();
              journal = null;
          }
      }
  
      private void recoverPendingLocalTransactions() throws IOException {
  
          if (checkForCorruption) {
              for (DataFile dataFile: journal.getFileMap().values()) {
                  if (!dataFile.getCorruptedBlocks().isEmpty()) {
                      LOG.error("Corrupt Transaction journal records found in db-{}.log at {}", dataFile.getDataFileId(),  dataFile.getCorruptedBlocks());
                      corruptJournalDetected.set(true);
                  }
              }
          }
          if (!corruptJournalDetected.get()) {
              Location location = null;
              try {
                  location = journal.getNextLocation(null);
                  while (location != null) {
                      process(location, load(location));
                      location = journal.getNextLocation(location);
                  }
              } catch (Exception oops) {
                  LOG.error("Corrupt journal record; unexpected exception on transaction journal replay of location:" + location, oops);
                  corruptJournalDetected.set(true);
              }
              pendingCommit.putAll(inflightTransactions);
              LOG.info("pending local transactions: " + pendingCommit.keySet());
          }
      }
  
      public JournalCommand<?> load(Location location) throws IOException {
          try(DataByteArrayInputStream is = new DataByteArrayInputStream(journal.read(location))) {
              byte readByte = is.readByte();
              KahaEntryType type = KahaEntryType.valueOf(readByte);
              if (type == null) {
                  throw new IOException("Could not load journal record. Invalid location: " + location);
              }
              JournalCommand<?> message = (JournalCommand<?>) type.createMessage();
              message.mergeFramed(is);
              return message;
          }
      }
  
      public void process(final Location location, JournalCommand<?> command) throws IOException {
          switch (command.type()) {
              case KAHA_PREPARE_COMMAND:
                  KahaPrepareCommand prepareCommand = (KahaPrepareCommand) command;
                  getTx(TransactionIdConversion.convert(prepareCommand.getTransactionInfo())).trackPrepareLocation(location);
                  break;
              case KAHA_COMMIT_COMMAND:
                  KahaCommitCommand commitCommand = (KahaCommitCommand) command;
                  removeTx(TransactionIdConversion.convert(commitCommand.getTransactionInfo()));
                  break;
              case KAHA_TRACE_COMMAND:
                  break;
              default:
                  throw new IOException("Unexpected command in transaction journal: " + command);
          }
      }
  
  
      @Override
      public synchronized void recover(final TransactionRecoveryListener listener) throws IOException {
  
          for (final PersistenceAdapter adapter : multiKahaDBPersistenceAdapter.adapters) {
              adapter.createTransactionStore().recover(new TransactionRecoveryListener() {
                  @Override
                  public void recover(XATransactionId xid, Message[] addedMessages, MessageAck[] acks) {
                      try {
                          getTx(xid).trackStore(adapter.createTransactionStore(), xid);
                      } catch (IOException e) {
                          LOG.error("Failed to access transaction store: " + adapter + " for prepared xa tid: " + xid, e);
                      }
                      listener.recover(xid, addedMessages, acks);
                  }
              });
          }
  
          boolean recoveryWorkPending = false;
          try {
              Broker broker = multiKahaDBPersistenceAdapter.getBrokerService().getBroker();
              // force completion of local xa
              for (TransactionId txid : broker.getPreparedTransactions(null)) {
                  if (multiKahaDBPersistenceAdapter.isLocalXid(txid)) {
                      recoveryWorkPending = true;
                      if (corruptJournalDetected.get()) {
                          // not having a record is meaningless once our tx store is corrupt; we need a heuristic decision
                          LOG.warn("Pending multi store local transaction {} requires manual heuristic outcome via JMX", txid);
                          logSomeContext(txid);
                      } else {
                          try {
                              if (pendingCommit.keySet().contains(txid)) {
                                  // we recorded the commit outcome, finish the job
                                  LOG.info("delivering pending commit outcome for tid: " + txid);
                                  broker.commitTransaction(null, txid, false);
                              } else {
                                  // we have not record an outcome, and would have reported a commit failure, so we must rollback
                                  LOG.info("delivering rollback outcome to store for tid: " + txid);
                                  broker.forgetTransaction(null, txid);
                              }
                              persistCompletion(txid);
                          } catch (Exception ex) {
                              LOG.error("failed to deliver pending outcome for tid: " + txid, ex);
                          }
                      }
                  }
              }
          } catch (Exception e) {
              LOG.error("failed to resolve pending local transactions", e);
          }
          // can we ignore corruption and resume
          if (corruptJournalDetected.get() && !recoveryWorkPending) {
              // move to new write file, gc will cleanup
              journal.rotateWriteFile();
              loaded();
              corruptJournalDetected.set(false);
              LOG.info("No heuristics outcome pending after corrupt tx store detection, auto resolving");
          }
      }
  
      private void logSomeContext(TransactionId txid) throws IOException {
          Tx tx = getTx(txid);
          if (tx != null) {
              for (TransactionStore store: tx.getStores()) {
                  for (PersistenceAdapter persistenceAdapter : multiKahaDBPersistenceAdapter.adapters) {
                      if (persistenceAdapter.createTransactionStore() == store) {
                          if (persistenceAdapter instanceof KahaDBPersistenceAdapter) {
                              LOG.warn("Heuristic data in: " + persistenceAdapter + ", "  + ((KahaDBPersistenceAdapter)persistenceAdapter).getStore().getPreparedTransaction(txid));
                          }
                      }
                  }
              }
          }
      }
  
      void addMessage(final TransactionStore transactionStore, ConnectionContext context, final MessageStore destination, final Message message)
              throws IOException {
          if (message.getTransactionId() != null) {
              getTx(message.getTransactionId()).trackStore(transactionStore);
          }
          destination.addMessage(context, message);
      }
  
      ListenableFuture<Object> asyncAddQueueMessage(final TransactionStore transactionStore, ConnectionContext context, final MessageStore destination, final Message message)
              throws IOException {
          if (message.getTransactionId() != null) {
              getTx(message.getTransactionId()).trackStore(transactionStore);
              destination.addMessage(context, message);
              return AbstractMessageStore.FUTURE;
          } else {
              return destination.asyncAddQueueMessage(context, message);
          }
      }
  
      ListenableFuture<Object> asyncAddTopicMessage(final TransactionStore transactionStore, ConnectionContext context, final MessageStore destination, final Message message)
              throws IOException {
  
          if (message.getTransactionId() != null) {
              getTx(message.getTransactionId()).trackStore(transactionStore);
              destination.addMessage(context, message);
              return AbstractMessageStore.FUTURE;
          } else {
              return destination.asyncAddTopicMessage(context, message);
          }
      }
  
      final void removeMessage(final TransactionStore transactionStore, ConnectionContext context, final MessageStore destination, final MessageAck ack)
              throws IOException {
          if (ack.getTransactionId() != null) {
              getTx(ack.getTransactionId()).trackStore(transactionStore);
          }
          destination.removeMessage(context, ack);
      }
  
      final void removeAsyncMessage(final TransactionStore transactionStore, ConnectionContext context, final MessageStore destination, final MessageAck ack)
              throws IOException {
          if (ack.getTransactionId() != null) {
              getTx(ack.getTransactionId()).trackStore(transactionStore);
          }
          destination.removeAsyncMessage(context, ack);
      }
  
      final void acknowledge(final TransactionStore transactionStore, ConnectionContext context, final TopicMessageStore destination,
                             final String clientId, final String subscriptionName,
                             final MessageId messageId, final MessageAck ack) throws IOException {
          if (ack.getTransactionId() != null) {
              getTx(ack.getTransactionId()).trackStore(transactionStore);
          }
          destination.acknowledge(context, clientId, subscriptionName, messageId, ack);
      }
  
  }
  