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.jdbc;
  
  import java.io.IOException;
  import java.sql.*;
  import java.util.LinkedList;
  import java.util.Map;
  import java.util.Properties;
  import java.util.concurrent.Executor;
  import java.util.concurrent.Executors;
  import java.util.concurrent.locks.Lock;
  import java.util.concurrent.locks.ReentrantReadWriteLock;
  
  import javax.sql.DataSource;
  
  import org.apache.activemq.util.IOExceptionSupport;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  /**
   * Helps keep track of the current transaction/JDBC connection.
   */
  public class TransactionContext {
  
      private static final Logger LOG = LoggerFactory.getLogger(TransactionContext.class);
  
      private final DataSource dataSource;
      private final JDBCPersistenceAdapter persistenceAdapter;
      private Connection connection;
      private boolean inTx;
      private PreparedStatement addMessageStatement;
      private PreparedStatement removedMessageStatement;
      private PreparedStatement updateLastAckStatement;
      // a cheap dirty level that we can live with    
      private int transactionIsolation = Connection.TRANSACTION_READ_UNCOMMITTED;
      private LinkedList<Runnable> completions = new LinkedList<Runnable>();
      private ReentrantReadWriteLock exclusiveConnectionLock = new ReentrantReadWriteLock();
      private int networkTimeout;
      private int queryTimeout;
  
      public TransactionContext(JDBCPersistenceAdapter persistenceAdapter, int networkTimeout, int queryTimeout) throws IOException {
          this.persistenceAdapter = persistenceAdapter;
          this.dataSource = persistenceAdapter.getDataSource();
          this.networkTimeout = networkTimeout;
          this.queryTimeout = queryTimeout;
      }
  
      public Connection getExclusiveConnection() throws IOException {
          return lockAndWrapped(exclusiveConnectionLock.writeLock());
      }
  
      public Connection getConnection() throws IOException {
          return lockAndWrapped(exclusiveConnectionLock.readLock());
      }
  
      private Connection lockAndWrapped(Lock toLock) throws IOException {
          if (connection == null) {
              toLock.lock();
              try {
                  connection = dataSource.getConnection();
                  if (networkTimeout > 0) {
                      connection.setNetworkTimeout(Executors.newSingleThreadExecutor(), networkTimeout);
                  }
                  if (persistenceAdapter.isChangeAutoCommitAllowed()) {
                      boolean autoCommit = !inTx;
                      if (connection.getAutoCommit() != autoCommit) {
                          LOG.trace("Setting auto commit to {} on connection {}", autoCommit, connection);
                          connection.setAutoCommit(autoCommit);
                      }
                  }
                  connection = new UnlockOnCloseConnection(connection, toLock);
              } catch (SQLException e) {
                  JDBCPersistenceAdapter.log("Could not get JDBC connection: ", e);
                  inTx = false;
                  try {
                      toLock.unlock();
                  } catch (IllegalMonitorStateException oops) {
                      LOG.error("Thread does not hold the context lock on close of:"  + connection, oops);
                  }
                  silentClose();
                  IOException ioe = IOExceptionSupport.create(e);
                  if (persistenceAdapter.getBrokerService() != null) {
                      persistenceAdapter.getBrokerService().handleIOException(ioe);
                  }
                  throw ioe;
              }
  
              try {
                  connection.setTransactionIsolation(transactionIsolation);
              } catch (Throwable e) {
                  // ignore
                  LOG.trace("Cannot set transaction isolation to " + transactionIsolation + " due " + e.getMessage()
                          + ". This exception is ignored.", e);
              }
          }
          return connection;
      }
  
      public void executeBatch() throws SQLException {
          try {
              executeBatch(addMessageStatement, "Failed add a message");
          } finally {
              addMessageStatement = null;
              try {
                  executeBatch(removedMessageStatement, "Failed to remove a message");
              } finally {
                  removedMessageStatement = null;
                  try {
                      executeBatch(updateLastAckStatement, "Failed to ack a message");
                  } finally {
                      updateLastAckStatement = null;
                  }
              }
          }
      }
  
      private void executeBatch(PreparedStatement p, String message) throws SQLException {
          if (p == null) {
              return;
          }
  
          try {
              int[] rc = p.executeBatch();
              for (int i = 0; i < rc.length; i++) {
                  int code = rc[i];
                  if (code < 0 && code != Statement.SUCCESS_NO_INFO) {
                      throw new SQLException(message + ". Response code: " + code);
                  }
              }
          } finally {
              try {
                  p.close();
              } catch (Throwable ignored) {}
          }
      }
  
      private void silentClose() {
          silentClosePreparedStatements();
          if (connection != null) {
              try {
                  connection.close();
              } catch (Throwable ignored) {}
              connection = null;
          }
      }
  
  
      public void close() throws IOException {
          if (!inTx) {
              try {
                  // can be null for topic ops that bypass the store via existing cursor state
                  if (connection != null) {
                      final boolean needsCommit = !connection.getAutoCommit();
                      executeBatch();
                      if (needsCommit) {
                          connection.commit();
                      }
                  }
              } catch (SQLException e) {
                  JDBCPersistenceAdapter.log("Error while closing connection: ", e);
                  IOException ioe = IOExceptionSupport.create(e);
                  persistenceAdapter.getBrokerService().handleIOException(ioe);
                  throw ioe;
              } finally {
                  silentClose();
                  for (Runnable completion: completions) {
                      completion.run();
                  }
                  completions.clear();
              }
          }
      }
  
      public void begin() throws IOException {
          if (inTx) {
              throw new IOException("Already started.");
          }
          inTx = true;
          connection = getConnection();
      }
  
      public void commit() throws IOException {
          if (!inTx) {
              throw new IOException("Not started.");
          }
          try {
              final boolean needsCommit = !connection.getAutoCommit();
              executeBatch();
              if (needsCommit) {
                  connection.commit();
              }
          } catch (SQLException e) {
              JDBCPersistenceAdapter.log("Commit failed: ", e);
              try {
                  doRollback();
              } catch (Exception ignored) {}
              IOException ioe = IOExceptionSupport.create(e);
              persistenceAdapter.getBrokerService().handleIOException(ioe);
              throw ioe;
          } finally {
              inTx = false;
              close();
          }
      }
  
      public void rollback() throws IOException {
          if (!inTx) {
              throw new IOException("Not started.");
          }
          try {
              doRollback();
          } catch (SQLException e) {
              JDBCPersistenceAdapter.log("Rollback failed: ", e);
              throw IOExceptionSupport.create(e);
          } finally {
              inTx = false;
              close();
          }
      }
  
      private PreparedStatement silentClosePreparedStatement(PreparedStatement preparedStatement) {
          if (preparedStatement != null) {
              try {
                  preparedStatement.close();
              } catch (Throwable ignored) {}
          }
          return null;
      }
  
      private void silentClosePreparedStatements() {
          addMessageStatement = silentClosePreparedStatement(addMessageStatement);
          removedMessageStatement = silentClosePreparedStatement(removedMessageStatement);
          updateLastAckStatement = silentClosePreparedStatement(updateLastAckStatement);
      }
  
      private void doRollback() throws SQLException {
          silentClosePreparedStatements();
          completions.clear();
          connection.rollback();
      }
  
      public PreparedStatement getAddMessageStatement() {
          return addMessageStatement;
      }
  
      public void setAddMessageStatement(PreparedStatement addMessageStatement) {
          this.addMessageStatement = addMessageStatement;
      }
  
      public PreparedStatement getUpdateLastAckStatement() {
          return updateLastAckStatement;
      }
  
      public void setUpdateLastAckStatement(PreparedStatement ackMessageStatement) {
          this.updateLastAckStatement = ackMessageStatement;
      }
  
      public PreparedStatement getRemovedMessageStatement() {
          return removedMessageStatement;
      }
  
      public void setRemovedMessageStatement(PreparedStatement removedMessageStatement) {
          this.removedMessageStatement = removedMessageStatement;
      }
      
      public void setTransactionIsolation(int transactionIsolation) {
          this.transactionIsolation = transactionIsolation;
      }
  
      public void onCompletion(Runnable runnable) {
          completions.add(runnable);
      }
  
      final private class UnlockOnCloseConnection implements Connection {
  
          private final Connection delegate;
          private final Lock lock;
  
          UnlockOnCloseConnection(Connection delegate, Lock toUnlockOnClose) {
              this.delegate = delegate;
              this.lock = toUnlockOnClose;
          }
  
          @Override
          public void close() throws SQLException {
              try {
                  delegate.close();
              } finally {
                  lock.unlock();
              }
          }
  
          // simple delegate for the  rest of the impl..
          @Override
          public Statement createStatement() throws SQLException {
              Statement statement = delegate.createStatement();
              if (queryTimeout > 0) {
                  statement.setQueryTimeout(queryTimeout);
              }
              return statement;
          }
  
          @Override
          public PreparedStatement prepareStatement(String sql) throws SQLException {
              PreparedStatement statement = delegate.prepareStatement(sql);
              if (queryTimeout > 0) {
                  statement.setQueryTimeout(queryTimeout);
              }
              return statement;
          }
  
          @Override
          public CallableStatement prepareCall(String sql) throws SQLException {
              CallableStatement statement = delegate.prepareCall(sql);
              if (queryTimeout > 0) {
                  statement.setQueryTimeout(queryTimeout);
              }
              return statement;
          }
  
          @Override
          public String nativeSQL(String sql) throws SQLException {
              return delegate.nativeSQL(sql);
          }
  
          @Override
          public void setAutoCommit(boolean autoCommit) throws SQLException {
              delegate.setAutoCommit(autoCommit);
          }
  
          @Override
          public boolean getAutoCommit() throws SQLException {
              return delegate.getAutoCommit();
          }
  
          @Override
          public void commit() throws SQLException {
              delegate.commit();
          }
  
          @Override
          public void rollback() throws SQLException {
              delegate.rollback();
          }
  
          @Override
          public boolean isClosed() throws SQLException {
              return delegate.isClosed();
          }
  
          @Override
          public DatabaseMetaData getMetaData() throws SQLException {
              return delegate.getMetaData();
          }
  
          @Override
          public void setReadOnly(boolean readOnly) throws SQLException {
              delegate.setReadOnly(readOnly);
          }
  
          @Override
          public boolean isReadOnly() throws SQLException {
              return delegate.isReadOnly();
          }
  
          @Override
          public void setCatalog(String catalog) throws SQLException {
              delegate.setCatalog(catalog);
          }
  
          @Override
          public String getCatalog() throws SQLException {
              return delegate.getCatalog();
          }
  
          @Override
          public void setTransactionIsolation(int level) throws SQLException {
              delegate.setTransactionIsolation(level);
          }
  
          @Override
          public int getTransactionIsolation() throws SQLException {
              return delegate.getTransactionIsolation();
          }
  
          @Override
          public SQLWarning getWarnings() throws SQLException {
              return delegate.getWarnings();
          }
  
          @Override
          public void clearWarnings() throws SQLException {
              delegate.clearWarnings();
          }
  
          @Override
          public Statement createStatement(int resultSetType, int resultSetConcurrency) throws SQLException {
              Statement statement = delegate.createStatement(resultSetType, resultSetConcurrency);
              if (queryTimeout > 0) {
                  statement.setQueryTimeout(queryTimeout);
              }
              return statement;
          }
  
          @Override
          public PreparedStatement prepareStatement(String sql, int resultSetType, int resultSetConcurrency) throws SQLException {
              PreparedStatement statement = delegate.prepareStatement(sql, resultSetType, resultSetConcurrency);
              if (queryTimeout > 0) {
                  statement.setQueryTimeout(queryTimeout);
              }
              return statement;
          }
  
          @Override
          public CallableStatement prepareCall(String sql, int resultSetType, int resultSetConcurrency) throws SQLException {
              CallableStatement statement = delegate.prepareCall(sql, resultSetType, resultSetConcurrency);
              if (queryTimeout > 0) {
                  statement.setQueryTimeout(queryTimeout);
              }
              return statement;
          }
  
          @Override
          public Map<String, Class<?>> getTypeMap() throws SQLException {
              return delegate.getTypeMap();
          }
  
          @Override
          public void setTypeMap(Map<String, Class<?>> map) throws SQLException {
              delegate.setTypeMap(map);
          }
  
          @Override
          public void setHoldability(int holdability) throws SQLException {
              delegate.setHoldability(holdability);
          }
  
          @Override
          public int getHoldability() throws SQLException {
              return delegate.getHoldability();
          }
  
          @Override
          public Savepoint setSavepoint() throws SQLException {
              return delegate.setSavepoint();
          }
  
          @Override
          public Savepoint setSavepoint(String name) throws SQLException {
              return delegate.setSavepoint(name);
          }
  
          @Override
          public void rollback(Savepoint savepoint) throws SQLException {
              delegate.rollback(savepoint);
          }
  
          @Override
          public void releaseSavepoint(Savepoint savepoint) throws SQLException {
              delegate.releaseSavepoint(savepoint);
          }
  
          @Override
          public Statement createStatement(int resultSetType, int resultSetConcurrency, int resultSetHoldability) throws SQLException {
              Statement statement = delegate.createStatement(resultSetType, resultSetConcurrency, resultSetHoldability);
              if (queryTimeout > 0) {
                  statement.setQueryTimeout(queryTimeout);
              }
              return statement;
          }
  
          @Override
          public PreparedStatement prepareStatement(String sql, int resultSetType, int resultSetConcurrency, int resultSetHoldability) throws SQLException {
              PreparedStatement statement = delegate.prepareStatement(sql, resultSetType, resultSetConcurrency, resultSetHoldability);
              if (queryTimeout > 0) {
                  statement.setQueryTimeout(queryTimeout);
              }
              return statement;
          }
  
          @Override
          public CallableStatement prepareCall(String sql, int resultSetType, int resultSetConcurrency, int resultSetHoldability) throws SQLException {
              CallableStatement statement = delegate.prepareCall(sql, resultSetType, resultSetConcurrency, resultSetHoldability);
              if (queryTimeout > 0) {
                  statement.setQueryTimeout(queryTimeout);
              }
              return statement;
          }
  
          @Override
          public PreparedStatement prepareStatement(String sql, int autoGeneratedKeys) throws SQLException {
              PreparedStatement statement = delegate.prepareStatement(sql, autoGeneratedKeys);
              if (queryTimeout > 0) {
                  statement.setQueryTimeout(queryTimeout);
              }
              return statement;
          }
  
          @Override
          public PreparedStatement prepareStatement(String sql, int[] columnIndexes) throws SQLException {
              PreparedStatement statement = delegate.prepareStatement(sql, columnIndexes);
              if (queryTimeout > 0) {
                  statement.setQueryTimeout(queryTimeout);
              }
              return statement;
          }
  
          @Override
          public PreparedStatement prepareStatement(String sql, String[] columnNames) throws SQLException {
              PreparedStatement statement = delegate.prepareStatement(sql, columnNames);
              if (queryTimeout > 0) {
                  statement.setQueryTimeout(queryTimeout);
              }
              return statement;
          }
  
          @Override
          public Clob createClob() throws SQLException {
              return delegate.createClob();
          }
  
          @Override
          public Blob createBlob() throws SQLException {
              return delegate.createBlob();
          }
  
          @Override
          public NClob createNClob() throws SQLException {
              return delegate.createNClob();
          }
  
          @Override
          public SQLXML createSQLXML() throws SQLException {
              return delegate.createSQLXML();
          }
  
          @Override
          public boolean isValid(int timeout) throws SQLException {
              return delegate.isValid(timeout);
          }
  
          @Override
          public void setClientInfo(String name, String value) throws SQLClientInfoException {
              delegate.setClientInfo(name, value);
          }
  
          @Override
          public void setClientInfo(Properties properties) throws SQLClientInfoException {
              delegate.setClientInfo(properties);
          }
  
          @Override
          public String getClientInfo(String name) throws SQLException {
              return delegate.getClientInfo(name);
          }
  
          @Override
          public Properties getClientInfo() throws SQLException {
              return delegate.getClientInfo();
          }
  
          @Override
          public Array createArrayOf(String typeName, Object[] elements) throws SQLException {
              return delegate.createArrayOf(typeName, elements);
          }
  
          @Override
          public Struct createStruct(String typeName, Object[] attributes) throws SQLException {
              return delegate.createStruct(typeName, attributes);
          }
  
          @Override
          public void setSchema(String schema) throws SQLException {
              delegate.setSchema(schema);
          }
  
          @Override
          public String getSchema() throws SQLException {
              return delegate.getSchema();
          }
  
          @Override
          public void abort(Executor executor) throws SQLException {
              delegate.abort(executor);
          }
  
          @Override
          public void setNetworkTimeout(Executor executor, int milliseconds) throws SQLException {
              delegate.setNetworkTimeout(executor, milliseconds);
          }
  
          @Override
          public int getNetworkTimeout() throws SQLException {
              return delegate.getNetworkTimeout();
          }
  
          @Override
          public <T> T unwrap(Class<T> iface) throws SQLException {
              return delegate.unwrap(iface);
          }
  
          @Override
          public boolean isWrapperFor(Class<?> iface) throws SQLException {
              return delegate.isWrapperFor(iface);
          }
      }
  }
  