sentry-core/sentry-core-common/src/main/java/org/apache/sentry/core/common/transport/SentryTransportPool.java - sentry - Git at Google

 /*
  * 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
  * <p>
  * http://www.apache.org/licenses/LICENSE-2.0
  * <p>
  * 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.sentry.core.common.transport;

 import com.google.common.base.Preconditions;
 import com.google.common.net.HostAndPort;
 import org.apache.commons.pool2.BaseKeyedPooledObjectFactory;
 import org.apache.commons.pool2.KeyedObjectPool;
 import org.apache.commons.pool2.PooledObject;
 import org.apache.commons.pool2.impl.DefaultPooledObject;
 import org.apache.commons.pool2.impl.GenericKeyedObjectPool;
 import org.apache.commons.pool2.impl.GenericKeyedObjectPoolConfig;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.sentry.core.common.utils.ThriftUtil;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;

 import javax.annotation.concurrent.ThreadSafe;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 import java.util.concurrent.atomic.AtomicBoolean;

 /**
  * Pool of transport connections to Sentry servers.
  * The pool caches open connections to multiple Sentry servers,
  * specified in the configuration.
  *
  * When transport pooling is disabled in configuration,
  * creates transports directly and doesn't cache connections.
  */
 @ThreadSafe
 public final class SentryTransportPool implements AutoCloseable {
   private static final Logger LOGGER = LoggerFactory.getLogger(SentryTransportPool.class);

   // Used for logging to identify pool instances. This is only useful for test debugging
   // so we do not preserve thread safety for this field.
   private static int poolId = 0;
   private final int id;

   // True if using Object pool
   private final boolean isPoolEnabled;

   // Load balance between servers if true
   private final boolean doLoadBalancing;

   // List of all known servers
   private final ArrayList<HostAndPort> endpoints;

   // Transport pool which keeps connected transports
   private final KeyedObjectPool<HostAndPort, TTransportWrapper> pool;
   // Source of connected transports
   private final TransportFactory transportFactory;

   // Set when we are closed
   private final AtomicBoolean closed = new AtomicBoolean();

   /**
    * Configure transport pool.
    * <p>
    * The pool accepts the following configuration:
    * <ul>
    *   <li>Maximum total number of objects in the pool</li>
    *   <li>Minimum number of idle objects</li>
    *   <li>Maximum number of idle objects</li>
    *   <li>Minimum time before the object is evicted</li>
    *   <li>Interval between evictions</li>
    * </ul>
    * @param conf Configuration
    * @param transportConfig Configuration interface
    * @param transportFactory Transport factory used to produce transports
    */
   public SentryTransportPool(Configuration conf,
                              SentryClientTransportConfigInterface transportConfig,
                              TransportFactory transportFactory) {

     // This isn't thread-safe, but we don't care - it is only used
     // for debugging when running tests - normal apps use a single pool
     poolId++;
     id = poolId;

     this.transportFactory = transportFactory;
     doLoadBalancing = transportConfig.isLoadBalancingEnabled(conf);
     isPoolEnabled = transportConfig.isTransportPoolEnabled(conf);

     // Get list of server addresses
     String hostsAndPortsStr = transportConfig.getSentryServerRpcAddress(conf);
     int serverPort = transportConfig.getServerRpcPort(conf);
     LOGGER.info("Creating pool for {} with default port {}",
             hostsAndPortsStr, serverPort);
     String[] hostsAndPortsStrArr = hostsAndPortsStr.split(",");
     Preconditions.checkArgument(hostsAndPortsStrArr.length > 0,
             "At least one server should be specified");

     endpoints = new ArrayList<>(hostsAndPortsStrArr.length);
     for(String addr: hostsAndPortsStrArr) {
       HostAndPort endpoint = ThriftUtil.parseAddress(addr, serverPort);
       LOGGER.info("Adding endpoint {}", endpoint);
       endpoints.add(endpoint);
     }

     if (!isPoolEnabled) {
       pool = null;
       LOGGER.info("Connection pooling is disabled");
       return;
     }

     LOGGER.info("Connection pooling is enabled");
     // Set pool configuration based on Configuration settings
     GenericKeyedObjectPoolConfig poolConfig = new GenericKeyedObjectPoolConfig();

     // Don't limit maximum number of objects in the pool
     poolConfig.setMaxTotal(-1);

     poolConfig.setMinIdlePerKey(transportConfig.getPoolMinIdle(conf));
     poolConfig.setMaxIdlePerKey(transportConfig.getPoolMaxIdle(conf));

     // Do not block when pool is exhausted, throw exception instead
     poolConfig.setBlockWhenExhausted(false);
     poolConfig.setTestOnReturn(true);

     // No limit for total objects in the pool
     poolConfig.setMaxTotalPerKey(transportConfig.getPoolMaxTotal(conf));
     poolConfig.setMinEvictableIdleTimeMillis(transportConfig.getMinEvictableTimeSec(conf));
     poolConfig.setTimeBetweenEvictionRunsMillis(transportConfig.getTimeBetweenEvictionRunsSec(conf));

     // Create object pool
     pool = new GenericKeyedObjectPool<>(new PoolFactory(this.transportFactory, id),
             poolConfig);
   }

   /**
    * Get an open transport instance.
    * The instance can be connected to any of the available servers.
    * We are trying to randomly load-balance between servers (unless it is
    * disabled in configuration).
    *
    * @return connected transport
    * @throws Exception if connection tto both servers fails
    */
   public TTransportWrapper getTransport() throws Exception {
     List<HostAndPort> servers;
     // If we are doing load balancing and there is more then one server,
     // shuffle them before obtaining connection
     if (doLoadBalancing && (endpoints.size() > 1)) {
       servers = new ArrayList<>(endpoints);
       Collections.shuffle(servers);
     } else {
       servers = endpoints;
     }

     // Try to get a connection from one of the pools.
     Exception failure = null;
     boolean ignoreEmptyPool = true;
     for (int attempt = 0; attempt < 2; attempt++) {
       // First only attempt to borrow from pools which have some idle connections
       // If this fails, try with all pools
       for (HostAndPort addr : servers) {
         if (isPoolEnabled && ignoreEmptyPool && (pool.getNumIdle(addr) == 0)) {
           LOGGER.debug("Ignoring empty pool {}", addr);
           ignoreEmptyPool = false;
           continue;
         }
         try {
           TTransportWrapper transport =
               isPoolEnabled ?
                   pool.borrowObject(addr) :
                   transportFactory.getTransport(addr);
           LOGGER.debug("[{}] obtained transport {}", id, transport);
           if (LOGGER.isDebugEnabled() && isPoolEnabled) {
             LOGGER.debug("Currently {} active connections, {} idle connections",
                 pool.getNumActive(), pool.getNumIdle());
           }
           return transport;
         } catch (IllegalStateException e) {
           // Should not happen
           LOGGER.error("Unexpected error from pool {}", id, e);
           failure = e;
         } catch (Exception e) {
           LOGGER.error("Failed to obtain transport for {}: {}",
               addr, e.getMessage());
           failure = e;
         }
       }
       ignoreEmptyPool = false;
     }
     // Failed to borrow connect to any endpoint
     assert failure != null;
     throw failure;
   }

   /**
    * Return transport to the pool
    * @param transport Open transport
    */
   public void returnTransport(TTransportWrapper transport) {
     if (closed.get()) {
       LOGGER.debug("Returned {} to closed pool", transport);
       transport.close();
       return;
     }
     try {
       if (isPoolEnabled) {
         LOGGER.debug("[{}] returning {}", id, transport);
         pool.returnObject(transport.getAddress(), transport);
       } else {
         LOGGER.debug("Closing {}", transport);
         transport.close();
       }
     } catch (Exception e) {
       LOGGER.error("Failed to return {}", transport, e);
     }
   }

   public void invalidateTransport(TTransportWrapper transport) {
     if (closed.get()) {
       LOGGER.debug("invalidated {} for closed pool", transport);
       transport.close();
       return;
     }
     try {
       LOGGER.debug("[{}] Invalidating address {}", id, transport);
       if (!isPoolEnabled) {
         transport.close();
       } else {
         pool.invalidateObject(transport.getAddress(), transport);
         // Invalidate the whole pool associated with the given address
         // It is a bit brutal since a single bad connection may
         // cause an invalidation, but otherwise we may have a lot of bad
         // connections in the pool and try to return them.
         pool.clear(transport.getAddress());
       }
     } catch (Exception e) {
       LOGGER.error("Failed to invalidate {}", transport, e);
     }
   }

   @Override
   public void close() throws Exception {
     if (closed.get()) {
       // already closed
       return;
     }
     LOGGER.debug("[{}] closing", id);
     if (pool != null) {
       LOGGER.debug("Closing pool of {}/{} endpoints",
               pool.getNumIdle(), pool.getNumActive());
       pool.close();
     }
   }

   /**
    * Factory that creates and destroys pool objects
    */
   private static final class PoolFactory
           extends BaseKeyedPooledObjectFactory<HostAndPort, TTransportWrapper> {
     private final TransportFactory transportFactory;
     private final int id;

     /**
      * Create a pool factory associated with the given transport factory
      * @param transportFactory - factory producing transports
      * @param id pool id (for debugging)
      */
     private PoolFactory(TransportFactory transportFactory, int id) {
       this.transportFactory = transportFactory;
       this.id = id;
     }

     @Override
     public boolean validateObject(HostAndPort key, PooledObject<TTransportWrapper> p) {
       TTransportWrapper transport = p.getObject();
       if (transport == null) {
         LOGGER.error("No transport to validate");
         return false;
       }
       if (transport.getAddress() != key) {
         LOGGER.error("Invalid endpoint {}: does not match {}", transport, key);
         return false;
       }
       return true;
     }

     @Override
     public TTransportWrapper create(HostAndPort key) throws Exception {
       TTransportWrapper transportWrapper = transportFactory.getTransport(key);
       LOGGER.debug("[{}] created {}", id, transportWrapper);
       return transportWrapper;
     }

     @Override
     public void destroyObject(HostAndPort key, PooledObject<TTransportWrapper> p) throws Exception {
       TTransportWrapper transport = p.getObject();
       if (transport != null) {
         LOGGER.debug("[{}] Destroying endpoint {}", id, transport);
         try {
           transport.close();
         } catch (RuntimeException e) {
           LOGGER.error("fail to destroy endpoint {}", transport, e);
         }
       }
     }

     @Override
     public PooledObject<TTransportWrapper> wrap(TTransportWrapper value) {
       return new DefaultPooledObject<>(value);
     }
   }

 }
	/*
	* 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
	* <p>
	* http://www.apache.org/licenses/LICENSE-2.0
	* <p>
	* 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.sentry.core.common.transport;

	import com.google.common.base.Preconditions;
	import com.google.common.net.HostAndPort;
	import org.apache.commons.pool2.BaseKeyedPooledObjectFactory;
	import org.apache.commons.pool2.KeyedObjectPool;
	import org.apache.commons.pool2.PooledObject;
	import org.apache.commons.pool2.impl.DefaultPooledObject;
	import org.apache.commons.pool2.impl.GenericKeyedObjectPool;
	import org.apache.commons.pool2.impl.GenericKeyedObjectPoolConfig;
	import org.apache.hadoop.conf.Configuration;
	import org.apache.sentry.core.common.utils.ThriftUtil;
	import org.slf4j.Logger;
	import org.slf4j.LoggerFactory;

	import javax.annotation.concurrent.ThreadSafe;
	import java.util.ArrayList;
	import java.util.Collections;
	import java.util.List;
	import java.util.concurrent.atomic.AtomicBoolean;

	/**
	* Pool of transport connections to Sentry servers.
	* The pool caches open connections to multiple Sentry servers,
	* specified in the configuration.
	*
	* When transport pooling is disabled in configuration,
	* creates transports directly and doesn't cache connections.
	*/
	@ThreadSafe
	public final class SentryTransportPool implements AutoCloseable {
	private static final Logger LOGGER = LoggerFactory.getLogger(SentryTransportPool.class);

	// Used for logging to identify pool instances. This is only useful for test debugging
	// so we do not preserve thread safety for this field.
	private static int poolId = 0;
	private final int id;

	// True if using Object pool
	private final boolean isPoolEnabled;

	// Load balance between servers if true
	private final boolean doLoadBalancing;

	// List of all known servers
	private final ArrayList<HostAndPort> endpoints;

	// Transport pool which keeps connected transports
	private final KeyedObjectPool<HostAndPort, TTransportWrapper> pool;
	// Source of connected transports
	private final TransportFactory transportFactory;

	// Set when we are closed
	private final AtomicBoolean closed = new AtomicBoolean();

	/**
	* Configure transport pool.
	* <p>
	* The pool accepts the following configuration:
	* <ul>
	* <li>Maximum total number of objects in the pool</li>
	* <li>Minimum number of idle objects</li>
	* <li>Maximum number of idle objects</li>
	* <li>Minimum time before the object is evicted</li>
	* <li>Interval between evictions</li>
	* </ul>
	* @param conf Configuration
	* @param transportConfig Configuration interface
	* @param transportFactory Transport factory used to produce transports
	*/
	public SentryTransportPool(Configuration conf,
	SentryClientTransportConfigInterface transportConfig,
	TransportFactory transportFactory) {

	// This isn't thread-safe, but we don't care - it is only used
	// for debugging when running tests - normal apps use a single pool
	poolId++;
	id = poolId;

	this.transportFactory = transportFactory;
	doLoadBalancing = transportConfig.isLoadBalancingEnabled(conf);
	isPoolEnabled = transportConfig.isTransportPoolEnabled(conf);

	// Get list of server addresses
	String hostsAndPortsStr = transportConfig.getSentryServerRpcAddress(conf);
	int serverPort = transportConfig.getServerRpcPort(conf);
	LOGGER.info("Creating pool for {} with default port {}",
	hostsAndPortsStr, serverPort);
	String[] hostsAndPortsStrArr = hostsAndPortsStr.split(",");
	Preconditions.checkArgument(hostsAndPortsStrArr.length > 0,
	"At least one server should be specified");

	endpoints = new ArrayList<>(hostsAndPortsStrArr.length);
	for(String addr: hostsAndPortsStrArr) {
	HostAndPort endpoint = ThriftUtil.parseAddress(addr, serverPort);
	LOGGER.info("Adding endpoint {}", endpoint);
	endpoints.add(endpoint);
	}

	if (!isPoolEnabled) {
	pool = null;
	LOGGER.info("Connection pooling is disabled");
	return;
	}

	LOGGER.info("Connection pooling is enabled");
	// Set pool configuration based on Configuration settings
	GenericKeyedObjectPoolConfig poolConfig = new GenericKeyedObjectPoolConfig();

	// Don't limit maximum number of objects in the pool
	poolConfig.setMaxTotal(-1);

	poolConfig.setMinIdlePerKey(transportConfig.getPoolMinIdle(conf));
	poolConfig.setMaxIdlePerKey(transportConfig.getPoolMaxIdle(conf));

	// Do not block when pool is exhausted, throw exception instead
	poolConfig.setBlockWhenExhausted(false);
	poolConfig.setTestOnReturn(true);

	// No limit for total objects in the pool
	poolConfig.setMaxTotalPerKey(transportConfig.getPoolMaxTotal(conf));
	poolConfig.setMinEvictableIdleTimeMillis(transportConfig.getMinEvictableTimeSec(conf));
	poolConfig.setTimeBetweenEvictionRunsMillis(transportConfig.getTimeBetweenEvictionRunsSec(conf));

	// Create object pool
	pool = new GenericKeyedObjectPool<>(new PoolFactory(this.transportFactory, id),
	poolConfig);
	}

	/**
	* Get an open transport instance.
	* The instance can be connected to any of the available servers.
	* We are trying to randomly load-balance between servers (unless it is
	* disabled in configuration).
	*
	* @return connected transport
	* @throws Exception if connection tto both servers fails
	*/
	public TTransportWrapper getTransport() throws Exception {
	List<HostAndPort> servers;
	// If we are doing load balancing and there is more then one server,
	// shuffle them before obtaining connection
	if (doLoadBalancing && (endpoints.size() > 1)) {
	servers = new ArrayList<>(endpoints);
	Collections.shuffle(servers);
	} else {
	servers = endpoints;
	}

	// Try to get a connection from one of the pools.
	Exception failure = null;
	boolean ignoreEmptyPool = true;
	for (int attempt = 0; attempt < 2; attempt++) {
	// First only attempt to borrow from pools which have some idle connections
	// If this fails, try with all pools
	for (HostAndPort addr : servers) {
	if (isPoolEnabled && ignoreEmptyPool && (pool.getNumIdle(addr) == 0)) {
	LOGGER.debug("Ignoring empty pool {}", addr);
	ignoreEmptyPool = false;
	continue;
	}
	try {
	TTransportWrapper transport =
	isPoolEnabled ?
	pool.borrowObject(addr) :
	transportFactory.getTransport(addr);
	LOGGER.debug("[{}] obtained transport {}", id, transport);
	if (LOGGER.isDebugEnabled() && isPoolEnabled) {
	LOGGER.debug("Currently {} active connections, {} idle connections",
	pool.getNumActive(), pool.getNumIdle());
	}
	return transport;
	} catch (IllegalStateException e) {
	// Should not happen
	LOGGER.error("Unexpected error from pool {}", id, e);
	failure = e;
	} catch (Exception e) {
	LOGGER.error("Failed to obtain transport for {}: {}",
	addr, e.getMessage());
	failure = e;
	}
	}
	ignoreEmptyPool = false;
	}
	// Failed to borrow connect to any endpoint
	assert failure != null;
	throw failure;
	}

	/**
	* Return transport to the pool
	* @param transport Open transport
	*/
	public void returnTransport(TTransportWrapper transport) {
	if (closed.get()) {
	LOGGER.debug("Returned {} to closed pool", transport);
	transport.close();
	return;
	}
	try {
	if (isPoolEnabled) {
	LOGGER.debug("[{}] returning {}", id, transport);
	pool.returnObject(transport.getAddress(), transport);
	} else {
	LOGGER.debug("Closing {}", transport);
	transport.close();
	}
	} catch (Exception e) {
	LOGGER.error("Failed to return {}", transport, e);
	}
	}

	public void invalidateTransport(TTransportWrapper transport) {
	if (closed.get()) {
	LOGGER.debug("invalidated {} for closed pool", transport);
	transport.close();
	return;
	}
	try {
	LOGGER.debug("[{}] Invalidating address {}", id, transport);
	if (!isPoolEnabled) {
	transport.close();
	} else {
	pool.invalidateObject(transport.getAddress(), transport);
	// Invalidate the whole pool associated with the given address
	// It is a bit brutal since a single bad connection may
	// cause an invalidation, but otherwise we may have a lot of bad
	// connections in the pool and try to return them.
	pool.clear(transport.getAddress());
	}
	} catch (Exception e) {
	LOGGER.error("Failed to invalidate {}", transport, e);
	}
	}

	@Override
	public void close() throws Exception {
	if (closed.get()) {
	// already closed
	return;
	}
	LOGGER.debug("[{}] closing", id);
	if (pool != null) {
	LOGGER.debug("Closing pool of {}/{} endpoints",
	pool.getNumIdle(), pool.getNumActive());
	pool.close();
	}
	}

	/**
	* Factory that creates and destroys pool objects
	*/
	private static final class PoolFactory
	extends BaseKeyedPooledObjectFactory<HostAndPort, TTransportWrapper> {
	private final TransportFactory transportFactory;
	private final int id;

	/**
	* Create a pool factory associated with the given transport factory
	* @param transportFactory - factory producing transports
	* @param id pool id (for debugging)
	*/
	private PoolFactory(TransportFactory transportFactory, int id) {
	this.transportFactory = transportFactory;
	this.id = id;
	}

	@Override
	public boolean validateObject(HostAndPort key, PooledObject<TTransportWrapper> p) {
	TTransportWrapper transport = p.getObject();
	if (transport == null) {
	LOGGER.error("No transport to validate");
	return false;
	}
	if (transport.getAddress() != key) {
	LOGGER.error("Invalid endpoint {}: does not match {}", transport, key);
	return false;
	}
	return true;
	}

	@Override
	public TTransportWrapper create(HostAndPort key) throws Exception {
	TTransportWrapper transportWrapper = transportFactory.getTransport(key);
	LOGGER.debug("[{}] created {}", id, transportWrapper);
	return transportWrapper;
	}

	@Override
	public void destroyObject(HostAndPort key, PooledObject<TTransportWrapper> p) throws Exception {
	TTransportWrapper transport = p.getObject();
	if (transport != null) {
	LOGGER.debug("[{}] Destroying endpoint {}", id, transport);
	try {
	transport.close();
	} catch (RuntimeException e) {
	LOGGER.error("fail to destroy endpoint {}", transport, e);
	}
	}
	}

	@Override
	public PooledObject<TTransportWrapper> wrap(TTransportWrapper value) {
	return new DefaultPooledObject<>(value);
	}
	}

	}