Separate error callbacks for master and sat

2 files changed, 53 insertions, 27 deletions

diff --git a/src/main/java/org/openslx/thrifthelper/ThriftHandler.java b/src/main/java/org/openslx/thrifthelper/ThriftHandler.java
index 84cd94d..c645dd8 100644
--- a/src/main/java/org/openslx/thrifthelper/ThriftHandler.java
+++ b/src/main/java/org/openslx/thrifthelper/ThriftHandler.java
@@ -11,9 +11,10 @@ import java.util.Set;
 
 import org.apache.log4j.Logger;
 import org.apache.thrift.TException;
+import org.apache.thrift.TServiceClient;
 import org.apache.thrift.transport.TTransportException;
 
-class ThriftHandler<T extends Object> implements InvocationHandler
+class ThriftHandler<T extends TServiceClient> implements InvocationHandler
 {
 
 	private final static Logger LOGGER = Logger.getLogger( ThriftHandler.class );
@@ -26,34 +27,31 @@ class ThriftHandler<T extends Object> implements InvocationHandler
 	}
 
 	private final Deque<T> pool = new ArrayDeque<>();
+
 	private final EventCallback<T> callback;
 
+	private final Set<String> thriftMethods;
+
 	public ThriftHandler( final Class<T> clazz, EventCallback<T> cb )
 	{
 		callback = cb;
-		thriftMethods = Collections.unmodifiableSet( new HashSet<String>() {
-			private static final long serialVersionUID = 8983506538154055231L;
-			{
-				Method[] methods = clazz.getMethods();
-				for ( int i = 0; i < methods.length; i++ ) {
-					boolean thrift = false;
-					Class<?>[] type = methods[i].getExceptionTypes();
-					for ( int e = 0; e < type.length; e++ ) {
-						if ( TException.class.isAssignableFrom( type[e] ) )
-							thrift = true;
-
-					}
-					String name = methods[i].getName();
-					if ( thrift && !name.startsWith( "send_" ) && !name.startsWith( "recv_" ) ) {
-						add( name );
-					}
-				}
+		Set<String> tmpset = new HashSet<String>();
+		Method[] methods = clazz.getMethods();
+		for ( int i = 0; i < methods.length; i++ ) {
+			boolean thrift = false;
+			Class<?>[] type = methods[i].getExceptionTypes();
+			for ( int e = 0; e < type.length; e++ ) {
+				if ( TException.class.isAssignableFrom( type[e] ) )
+					thrift = true;
 			}
-		} );
+			String name = methods[i].getName();
+			if ( thrift && !name.startsWith( "send_" ) && !name.startsWith( "recv_" ) ) {
+				tmpset.add( name );
+			}
+		}
+		thriftMethods = Collections.unmodifiableSet( tmpset );
 	}
 
-	private final Set<String> thriftMethods;
-
 	@Override
 	public Object invoke( Object tproxy, Method method, Object[] args ) throws Throwable
 	{
@@ -87,6 +85,7 @@ class ThriftHandler<T extends Object> implements InvocationHandler
 						if ( cause != null && ! ( cause instanceof TTransportException ) ) {
 							throw cause;
 						}
+						freeClient(client);
 						client = null;
 						if ( cause == null )
 							cause = e;
@@ -107,6 +106,18 @@ class ThriftHandler<T extends Object> implements InvocationHandler
 		}
 	}
 
+	private void freeClient(T client) {
+		try {
+			client.getInputProtocol().getTransport().close();
+		} catch (Exception e) {
+		}
+		try {
+			client.getOutputProtocol().getTransport().close();
+		} catch (Exception e) {
+		}
+		
+	}
+
 	private T getClient()
 	{
 		T client;
diff --git a/src/main/java/org/openslx/thrifthelper/ThriftManager.java b/src/main/java/org/openslx/thrifthelper/ThriftManager.java
index b132a9e..56b99b3 100644
--- a/src/main/java/org/openslx/thrifthelper/ThriftManager.java
+++ b/src/main/java/org/openslx/thrifthelper/ThriftManager.java
@@ -32,7 +32,9 @@ public class ThriftManager
 		public boolean thriftError( int failCount, String method, Throwable t );
 	}
 
-	private static ErrorCallback errorCallback = null;
+	private static ErrorCallback masterErrorCallback = null;
+	
+	private static ErrorCallback satelliteErrorCallback = null;
 
 	/**
 	 * Private members for master connection information
@@ -70,7 +72,7 @@ public class ThriftManager
 						@Override
 						public boolean error( int failCount, String method, Throwable t )
 						{
-							return errorCallback != null && errorCallback.thriftError( failCount, method, t );
+							return satelliteErrorCallback != null && satelliteErrorCallback.thriftError( failCount, method, t );
 						}
 					} ) );
 
@@ -113,7 +115,7 @@ public class ThriftManager
 						public boolean error( int failCount, String method, Throwable t )
 						{
 							synchronized ( LOGGER ) {
-								return errorCallback != null && errorCallback.thriftError( failCount, method, t );
+								return masterErrorCallback != null && masterErrorCallback.thriftError( failCount, method, t );
 							}
 						}
 					} ) );
@@ -182,14 +184,27 @@ public class ThriftManager
 
 	/**
 	 * Set the callback class for errors that occur on one of the
-	 * thrift connections.
+	 * thrift connections to the master server.
 	 * 
 	 * @param cb
 	 */
-	public static void setErrorCallback( ErrorCallback cb )
+	public static void setMasterErrorCallback( ErrorCallback cb )
 	{
 		synchronized ( LOGGER ) {
-			errorCallback = cb;
+			masterErrorCallback = cb;
+		}
+	}
+
+	/**
+	 * Set the callback class for errors that occur on one of the
+	 * thrift connections to the satellite server.
+	 * 
+	 * @param cb
+	 */
+	public static void setSatelliteErrorCallback( ErrorCallback cb )
+	{
+		synchronized ( LOGGER ) {
+			satelliteErrorCallback = cb;
 		}
 	}