/**
 * XaResourceCapableTransactionImpl.
 *
 * @author Galder Zamarreño
 * @since 3.5
 */
public class XaTransactionImpl implements Transaction {
  private static final InfinispanMessageLogger log =
      InfinispanMessageLogger.Provider.getLog(XaTransactionImpl.class);

  private int status;
  private LinkedList synchronizations;
  private Connection connection; // the only resource we care about is jdbc connection
  private final XaTransactionManagerImpl jtaTransactionManager;
  private List<XAResource> enlistedResources = new ArrayList<XAResource>();
  private Xid xid = new XaResourceCapableTransactionXid();
  private ConnectionProvider connectionProvider;

  public XaTransactionImpl(XaTransactionManagerImpl jtaTransactionManager) {
    this.jtaTransactionManager = jtaTransactionManager;
    this.status = Status.STATUS_ACTIVE;
  }

  public XaTransactionImpl(XaTransactionManagerImpl jtaTransactionManager, Xid xid) {
    this.jtaTransactionManager = jtaTransactionManager;
    this.status = Status.STATUS_ACTIVE;
    this.xid = xid;
  }

  public int getStatus() {
    return status;
  }

  public void commit()
      throws RollbackException, HeuristicMixedException, HeuristicRollbackException,
          IllegalStateException, SystemException {

    if (status == Status.STATUS_MARKED_ROLLBACK) {
      log.trace("on commit, status was marked for rollback-only");
      rollback();
    } else {
      status = Status.STATUS_PREPARING;

      if (synchronizations != null) {
        for (int i = 0; i < synchronizations.size(); i++) {
          Synchronization s = (Synchronization) synchronizations.get(i);
          s.beforeCompletion();
        }
      }

      if (!runXaResourcePrepare()) {
        status = Status.STATUS_ROLLING_BACK;
      } else {
        status = Status.STATUS_PREPARED;
      }

      status = Status.STATUS_COMMITTING;

      if (connection != null) {
        try {
          connection.commit();
          connectionProvider.closeConnection(connection);
          connection = null;
        } catch (SQLException sqle) {
          status = Status.STATUS_UNKNOWN;
          throw new SystemException();
        }
      }

      runXaResourceCommitTx();

      status = Status.STATUS_COMMITTED;

      if (synchronizations != null) {
        for (int i = 0; i < synchronizations.size(); i++) {
          Synchronization s = (Synchronization) synchronizations.get(i);
          s.afterCompletion(status);
        }
      }

      // status = Status.STATUS_NO_TRANSACTION;
      jtaTransactionManager.endCurrent(this);
    }
  }

  public void rollback() throws IllegalStateException, SystemException {
    status = Status.STATUS_ROLLING_BACK;
    runXaResourceRollback();
    status = Status.STATUS_ROLLEDBACK;

    if (connection != null) {
      try {
        connection.rollback();
        connection.close();
      } catch (SQLException sqle) {
        status = Status.STATUS_UNKNOWN;
        throw new SystemException();
      }
    }

    if (synchronizations != null) {
      for (int i = 0; i < synchronizations.size(); i++) {
        Synchronization s = (Synchronization) synchronizations.get(i);
        if (s != null) s.afterCompletion(status);
      }
    }

    // status = Status.STATUS_NO_TRANSACTION;
    jtaTransactionManager.endCurrent(this);
  }

  public void setRollbackOnly() throws IllegalStateException, SystemException {
    status = Status.STATUS_MARKED_ROLLBACK;
  }

  public void registerSynchronization(Synchronization synchronization)
      throws RollbackException, IllegalStateException, SystemException {
    // todo : find the spec-allowable statuses during which synch can be registered...
    if (synchronizations == null) {
      synchronizations = new LinkedList();
    }
    synchronizations.add(synchronization);
  }

  public void enlistConnection(Connection connection, ConnectionProvider connectionProvider) {
    if (this.connection != null) {
      throw new IllegalStateException("Connection already registered");
    }
    this.connection = connection;
    this.connectionProvider = connectionProvider;
  }

  public Connection getEnlistedConnection() {
    return connection;
  }

  public boolean enlistResource(XAResource xaResource)
      throws RollbackException, IllegalStateException, SystemException {
    enlistedResources.add(new WrappedXaResource(xaResource));
    try {
      xaResource.start(xid, 0);
    } catch (XAException e) {
      log.error("Got an exception", e);
      throw new SystemException(e.getMessage());
    }
    return true;
  }

  public boolean delistResource(XAResource xaResource, int i)
      throws IllegalStateException, SystemException {
    throw new SystemException("not supported");
  }

  public Collection<XAResource> getEnlistedResources() {
    return enlistedResources;
  }

  private boolean runXaResourcePrepare() throws SystemException {
    Collection<XAResource> resources = getEnlistedResources();
    for (XAResource res : resources) {
      try {
        res.prepare(xid);
      } catch (XAException e) {
        log.trace("The resource wants to rollback!", e);
        return false;
      } catch (Throwable th) {
        log.error("Unexpected error from resource manager!", th);
        throw new SystemException(th.getMessage());
      }
    }
    return true;
  }

  private void runXaResourceRollback() {
    Collection<XAResource> resources = getEnlistedResources();
    for (XAResource res : resources) {
      try {
        res.rollback(xid);
      } catch (XAException e) {
        log.warn("Error while rolling back", e);
      }
    }
  }

  private boolean runXaResourceCommitTx() throws HeuristicMixedException {
    Collection<XAResource> resources = getEnlistedResources();
    for (XAResource res : resources) {
      try {
        res.commit(xid, false); // todo we only support one phase commit for now, change this!!!
      } catch (XAException e) {
        log.warn("exception while committing", e);
        throw new HeuristicMixedException(e.getMessage());
      }
    }
    return true;
  }

  private static class XaResourceCapableTransactionXid implements Xid {
    private static AtomicInteger txIdCounter = new AtomicInteger(0);
    private int id = txIdCounter.incrementAndGet();

    public int getFormatId() {
      return id;
    }

    public byte[] getGlobalTransactionId() {
      throw new IllegalStateException("TODO - please implement me!!!"); // todo implement!!!
    }

    public byte[] getBranchQualifier() {
      throw new IllegalStateException("TODO - please implement me!!!"); // todo implement!!!
    }

    @Override
    public String toString() {
      return getClass().getSimpleName() + "{" + "id=" + id + '}';
    }
  }

  private class WrappedXaResource implements XAResource {
    private final XAResource xaResource;
    private int prepareResult;

    public WrappedXaResource(XAResource xaResource) {
      this.xaResource = xaResource;
    }

    @Override
    public void commit(Xid xid, boolean b) throws XAException {
      // Commit only if not read only.
      if (prepareResult != XAResource.XA_RDONLY) xaResource.commit(xid, b);
      else log.tracef("Not committing {0} due to readonly.", xid);
    }

    @Override
    public void end(Xid xid, int i) throws XAException {
      xaResource.end(xid, i);
    }

    @Override
    public void forget(Xid xid) throws XAException {
      xaResource.forget(xid);
    }

    @Override
    public int getTransactionTimeout() throws XAException {
      return xaResource.getTransactionTimeout();
    }

    @Override
    public boolean isSameRM(XAResource xaResource) throws XAException {
      return xaResource.isSameRM(xaResource);
    }

    @Override
    public int prepare(Xid xid) throws XAException {
      prepareResult = xaResource.prepare(xid);
      return prepareResult;
    }

    @Override
    public Xid[] recover(int i) throws XAException {
      return xaResource.recover(i);
    }

    @Override
    public void rollback(Xid xid) throws XAException {
      xaResource.rollback(xid);
    }

    @Override
    public boolean setTransactionTimeout(int i) throws XAException {
      return xaResource.setTransactionTimeout(i);
    }

    @Override
    public void start(Xid xid, int i) throws XAException {
      xaResource.start(xid, i);
    }
  }
}
/**
 * Variant of SimpleJtaTransactionManagerImpl that doesn't use a VM-singleton, but rather a set of
 * impls keyed by a node id.
 *
 * <p>TODO: Merge with single node transaction manager as much as possible
 *
 * @author Brian Stansberry
 */
public class DualNodeJtaTransactionManagerImpl implements TransactionManager {

  private static final InfinispanMessageLogger log =
      InfinispanMessageLogger.Provider.getLog(DualNodeJtaTransactionManagerImpl.class);

  private static final Hashtable INSTANCES = new Hashtable();

  private ThreadLocal currentTransaction = new ThreadLocal();
  private String nodeId;

  public static synchronized DualNodeJtaTransactionManagerImpl getInstance(String nodeId) {
    DualNodeJtaTransactionManagerImpl tm =
        (DualNodeJtaTransactionManagerImpl) INSTANCES.get(nodeId);
    if (tm == null) {
      tm = new DualNodeJtaTransactionManagerImpl(nodeId);
      INSTANCES.put(nodeId, tm);
    }
    return tm;
  }

  public static synchronized void cleanupTransactions() {
    for (java.util.Iterator it = INSTANCES.values().iterator(); it.hasNext(); ) {
      TransactionManager tm = (TransactionManager) it.next();
      try {
        tm.suspend();
      } catch (Exception e) {
        log.error("Exception cleaning up TransactionManager " + tm);
      }
    }
  }

  public static synchronized void cleanupTransactionManagers() {
    INSTANCES.clear();
  }

  private DualNodeJtaTransactionManagerImpl(String nodeId) {
    this.nodeId = nodeId;
  }

  public int getStatus() throws SystemException {
    Transaction tx = getCurrentTransaction();
    return tx == null ? Status.STATUS_NO_TRANSACTION : tx.getStatus();
  }

  public Transaction getTransaction() throws SystemException {
    return (Transaction) currentTransaction.get();
  }

  public DualNodeJtaTransactionImpl getCurrentTransaction() {
    return (DualNodeJtaTransactionImpl) currentTransaction.get();
  }

  public void begin() throws NotSupportedException, SystemException {
    currentTransaction.set(new DualNodeJtaTransactionImpl(this));
  }

  public Transaction suspend() throws SystemException {
    DualNodeJtaTransactionImpl suspended = getCurrentTransaction();
    log.trace(
        nodeId + ": Suspending " + suspended + " for thread " + Thread.currentThread().getName());
    currentTransaction.set(null);
    return suspended;
  }

  public void resume(Transaction transaction)
      throws InvalidTransactionException, IllegalStateException, SystemException {
    currentTransaction.set(transaction);
    log.trace(
        nodeId + ": Resumed " + transaction + " for thread " + Thread.currentThread().getName());
  }

  public void commit()
      throws RollbackException, HeuristicMixedException, HeuristicRollbackException,
          SecurityException, IllegalStateException, SystemException {
    Transaction tx = getCurrentTransaction();
    if (tx == null) {
      throw new IllegalStateException("no current transaction to commit");
    }
    tx.commit();
  }

  public void rollback() throws IllegalStateException, SecurityException, SystemException {
    Transaction tx = getCurrentTransaction();
    if (tx == null) {
      throw new IllegalStateException("no current transaction");
    }
    tx.rollback();
  }

  public void setRollbackOnly() throws IllegalStateException, SystemException {
    Transaction tx = getCurrentTransaction();
    if (tx == null) {
      throw new IllegalStateException("no current transaction");
    }
    tx.setRollbackOnly();
  }

  public void setTransactionTimeout(int i) throws SystemException {}

  void endCurrent(DualNodeJtaTransactionImpl transaction) {
    if (transaction == currentTransaction.get()) {
      currentTransaction.set(null);
    }
  }

  @Override
  public String toString() {
    StringBuffer sb = new StringBuffer(getClass().getName());
    sb.append("[nodeId=");
    sb.append(nodeId);
    sb.append("]");
    return sb.toString();
  }
}
/**
 * Intercepts transactions in Infinispan, calling {@link
 * PutFromLoadValidator#beginInvalidatingKey(Object, Object)} beforeQuery locks are acquired (and
 * the entry is invalidated) and sends {@link EndInvalidationCommand} to release invalidation
 * throught {@link PutFromLoadValidator#endInvalidatingKey(Object, Object)} afterQuery the
 * transaction is committed.
 *
 * @author Radim Vansa &lt;[email protected]&gt;
 */
class TxPutFromLoadInterceptor extends BaseRpcInterceptor {
  private static final InfinispanMessageLogger log =
      InfinispanMessageLogger.Provider.getLog(TxPutFromLoadInterceptor.class);
  private PutFromLoadValidator putFromLoadValidator;
  private final String cacheName;
  private RpcManager rpcManager;
  private CacheCommandInitializer cacheCommandInitializer;
  private DataContainer dataContainer;

  public TxPutFromLoadInterceptor(PutFromLoadValidator putFromLoadValidator, String cacheName) {
    this.putFromLoadValidator = putFromLoadValidator;
    this.cacheName = cacheName;
  }

  @Inject
  public void injectDependencies(
      RpcManager rpcManager,
      CacheCommandInitializer cacheCommandInitializer,
      DataContainer dataContainer) {
    this.rpcManager = rpcManager;
    this.cacheCommandInitializer = cacheCommandInitializer;
    this.dataContainer = dataContainer;
  }

  // We need to intercept PrepareCommand, not InvalidateCommand since the interception takes
  // place beforeQuery EntryWrappingInterceptor and the PrepareCommand is multiplexed into
  // InvalidateCommands
  // as part of EntryWrappingInterceptor
  @Override
  public Object visitPrepareCommand(TxInvocationContext ctx, PrepareCommand command)
      throws Throwable {
    if (ctx.isOriginLocal()) {
      // We can't wait to commit phase to remove the entry locally (invalidations are processed in
      // 1pc
      // on remote nodes, so only local case matters here). The problem is that while the entry is
      // locked
      // reads still can take place and we can read outdated collection afterQuery reading updated
      // entity
      // owning this collection from DB; when this happens, the version lock on entity cannot
      // protect
      // us against concurrent modification of the collection. Therefore, we need to remove the
      // entry
      // here (even without lock!) and let possible update happen in commit phase.
      for (WriteCommand wc : command.getModifications()) {
        if (wc instanceof InvalidateCommand) {
          // ISPN-5605 InvalidateCommand does not correctly implement getAffectedKeys()
          for (Object key : ((InvalidateCommand) wc).getKeys()) {
            dataContainer.remove(key);
          }
        } else {
          for (Object key : wc.getAffectedKeys()) {
            dataContainer.remove(key);
          }
        }
      }
    } else {
      for (WriteCommand wc : command.getModifications()) {
        if (wc instanceof InvalidateCommand) {
          // ISPN-5605 InvalidateCommand does not correctly implement getAffectedKeys()
          for (Object key : ((InvalidateCommand) wc).getKeys()) {
            if (log.isTraceEnabled()) {
              log.tracef("Invalidating key %s with lock owner %s", key, ctx.getLockOwner());
            }
            putFromLoadValidator.beginInvalidatingKey(ctx.getLockOwner(), key);
          }
        } else {
          Set<Object> keys = wc.getAffectedKeys();
          if (log.isTraceEnabled()) {
            log.tracef("Invalidating keys %s with lock owner %s", keys, ctx.getLockOwner());
          }
          for (Object key : keys) {
            putFromLoadValidator.beginInvalidatingKey(ctx.getLockOwner(), key);
          }
        }
      }
    }
    return invokeNextInterceptor(ctx, command);
  }

  @Override
  public Object visitCommitCommand(TxInvocationContext ctx, CommitCommand command)
      throws Throwable {
    if (log.isTraceEnabled()) {
      log.tracef("Commit command received, end invalidation");
    }

    return endInvalidationAndInvokeNextInterceptor(ctx, command);
  }

  @Override
  public Object visitRollbackCommand(TxInvocationContext ctx, RollbackCommand command)
      throws Throwable {
    if (log.isTraceEnabled()) {
      log.tracef("Rollback command received, end invalidation");
    }

    return endInvalidationAndInvokeNextInterceptor(ctx, command);
  }

  protected Object endInvalidationAndInvokeNextInterceptor(
      TxInvocationContext ctx, VisitableCommand command) throws Throwable {
    try {
      if (ctx.isOriginLocal()) {
        // send async Commit
        Set<Object> affectedKeys = ctx.getAffectedKeys();

        if (log.isTraceEnabled()) {
          log.tracef("Sending end invalidation for keys %s asynchronously", affectedKeys);
        }

        if (!affectedKeys.isEmpty()) {
          EndInvalidationCommand commitCommand =
              cacheCommandInitializer.buildEndInvalidationCommand(
                  cacheName, affectedKeys.toArray(), ctx.getGlobalTransaction());
          rpcManager.invokeRemotely(
              null, commitCommand, rpcManager.getDefaultRpcOptions(false, DeliverOrder.NONE));
        }
      }
    } finally {
      return invokeNextInterceptor(ctx, command);
    }
  }
}