AbstractFunctionalMap.java

package org.infinispan.functional.impl;

import java.util.concurrent.CompletableFuture;

import javax.transaction.SystemException;
import javax.transaction.Transaction;
import javax.transaction.TransactionManager;

import org.infinispan.batch.BatchContainer;
import org.infinispan.commands.VisitableCommand;
import org.infinispan.commons.CacheException;
import org.infinispan.commons.api.functional.FunctionalMap;
import org.infinispan.commons.api.functional.Status;
import org.infinispan.util.concurrent.CompletableFutures;
import org.infinispan.util.logging.Log;
import org.infinispan.util.logging.LogFactory;
import org.infinispan.configuration.cache.Configuration;
import org.infinispan.context.InvocationContext;
import org.infinispan.context.impl.TxInvocationContext;

/**
 * Abstract functional map, providing implementations for some of the shared methods.
 *
 * @since 8.0
 */
abstract class AbstractFunctionalMap<K, V> implements FunctionalMap<K, V> {
   private static final Log log = LogFactory.getLog(FunctionalMap.class);

   protected final FunctionalMapImpl<K, V> fmap;
   protected final Params params;

   private final boolean transactional;
   private final boolean autoCommit;
   private final BatchContainer batchContainer;
   private final TransactionManager transactionManager;

   protected AbstractFunctionalMap(Params params, FunctionalMapImpl<K, V> fmap) {
      this.fmap = fmap;
      Configuration config = fmap.cache.getCacheConfiguration();
      transactional = config.transaction().transactionMode().isTransactional();
      autoCommit = config.transaction().autoCommit();
      transactionManager = transactional ? fmap.cache.getTransactionManager() : null;
      batchContainer = transactional && config.invocationBatching().enabled() ? fmap.cache.getBatchContainer() : null;
      this.params = params;
   }

   @Override
   public String getName() {
      return "";
   }

   @Override
   public Status getStatus() {
      return fmap.getStatus();
   }

   @Override
   public void close() throws Exception {
      fmap.close();
   }

   protected InvocationContext getInvocationContext(boolean isWrite, int keyCount) {
      InvocationContext invocationContext;
      boolean txInjected = false;
      if (transactional) {
         Transaction transaction = getOngoingTransaction();
         if (transaction == null && autoCommit && transactionManager != null) {
            try {
               transactionManager.begin();
            } catch (RuntimeException e) {
               throw e;
            } catch (Exception e) {
               throw new CacheException("Unable to begin implicit transaction.", e);
            }
            transaction = getOngoingTransaction();
            txInjected = true;
         }
         invocationContext = fmap.invCtxFactory.createInvocationContext(transaction, txInjected);
      } else {
         invocationContext = fmap.invCtxFactory.createInvocationContext(isWrite, keyCount);
      }
      return invocationContext;
   }

   private Transaction getOngoingTransaction() {
      try {
         Transaction transaction = null;
         if (transactionManager != null) {
            transaction = transactionManager.getTransaction();
            if (transaction == null && batchContainer != null) {
               transaction = batchContainer.getBatchTransaction();
            }
         }
         return transaction;
      } catch (SystemException e) {
         throw new CacheException("Unable to get transaction", e);
      }
   }

   protected <T> CompletableFuture<T> invokeAsync(InvocationContext ctx, VisitableCommand cmd) {
      CompletableFuture<T> cf;
      boolean isImplicitTx = ctx.isInTxScope() && ((TxInvocationContext) ctx).isImplicitTransaction();
      final Transaction implicitTransaction;
      try {
         // interceptors must not access thread-local transaction anyway
         if (isImplicitTx) {
            implicitTransaction = transactionManager.suspend();
            assert implicitTransaction != null;
         } else {
            implicitTransaction = null;
         }
         cf = (CompletableFuture<T>) fmap.chain.invokeAsync(ctx, cmd);
      } catch (SystemException e) {
         throw new CacheException("Cannot suspend implicit transaction", e);
      } catch (Throwable t) {
         if (isImplicitTx) {
            try {
               if (transactionManager != null) transactionManager.rollback();
            } catch (Throwable t2) {
               log.trace("Could not rollback", t2);//best effort
               t.addSuppressed(t2);
            }
         }
         throw t;
      }
      if (isImplicitTx) {
         return cf.handle((result, throwable) -> {
            if (throwable != null) {
               try {
                  implicitTransaction.rollback();
               } catch (SystemException e) {
                  log.trace("Could not rollback", e);
                  throwable.addSuppressed(e);
               }
               throw CompletableFutures.asCompletionException(throwable);
            }
            try {
               implicitTransaction.commit();
            } catch (Exception e) {
               log.couldNotCompleteInjectedTransaction(e);
               throw CompletableFutures.asCompletionException(e);
            }
            return result;
         });
      } else {
         return cf;
      }
   }
}