summaryrefslogtreecommitdiffstats
path: root/src/main/java/org/openslx/thrifthelper/ThriftHandler.java
blob: 37c2897bcbf0bd20dfe5036e37cca1ce83822d96 (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
package org.openslx.thrifthelper;

import java.lang.reflect.InvocationHandler;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.util.ArrayDeque;
import java.util.Collections;
import java.util.Deque;
import java.util.HashSet;
import java.util.Iterator;
import java.util.Set;

import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.apache.thrift.TException;
import org.apache.thrift.TServiceClient;
import org.apache.thrift.protocol.TProtocolException;
import org.apache.thrift.transport.TTransportException;
import org.openslx.thrifthelper.ThriftManager.ErrorCallback;
import org.openslx.util.QuickTimer;
import org.openslx.util.QuickTimer.Task;

class ThriftHandler<T extends TServiceClient> implements InvocationHandler
{

	private final static Logger LOGGER = LogManager.getLogger( ThriftHandler.class );
	
	/**
	 * How long a client/connection should be allowed to be idle before we get rid of it and open a fresh one
	 */
	private final static long MAX_IDLE_MS = 110_000;

	protected interface WantClientCallback<T>
	{
		public T getNewClient();
	}

	/** Pool of (potentially) good connections, waiting for reuse */
	private final Deque<TWrap<T>> pool = new ArrayDeque<>();
	/** Old connections to be released by cleanup task */
	private Deque<TWrap<T>> trash = new ArrayDeque<>();

	private final WantClientCallback<? extends T> clientFactory;

	private final ErrorCallback errorCallback;

	private final Set<String> thriftMethods;

	public ThriftHandler( final Class<? extends T> clazz, WantClientCallback<? extends T> clientFactory, ErrorCallback errCb )
	{
		this.errorCallback = errCb;
		this.clientFactory = clientFactory;
		Set<String> tmpset = new HashSet<String>();
		Method[] methods = clazz.getMethods();
		// Iterate over all methods of this class
		for ( int i = 0; i < methods.length; i++ ) {
			boolean thrift = false;
			// If a method throws some form of TException, consider it a potential method from the thrift interface
			Class<?>[] type = methods[i].getExceptionTypes();
			for ( int e = 0; e < type.length; e++ ) {
				if ( TException.class.isAssignableFrom( type[e] ) ) {
					thrift = true;
					break;
				}
			}
			String name = methods[i].getName();
			if ( thrift && !name.startsWith( "send_" ) && !name.startsWith( "recv_" ) ) {
				// Exclude the send_ and recv_ helpers
				tmpset.add( name );
			}
		}
		thriftMethods = Collections.unmodifiableSet( tmpset );
		// Periodically scan for old connections, in case the application idles for extended periods of time...
		QuickTimer.scheduleAtFixedDelay( new Task() {
			@Override
			public void fire()
			{
				Deque<TWrap<T>> newTrash = new ArrayDeque<>();
				Deque<TWrap<T>> list;
				long now = System.currentTimeMillis();
				synchronized ( pool ) {
					list = trash;
					trash = newTrash;
					for ( Iterator<TWrap<T>> it = pool.iterator(); it.hasNext(); ) {
						TWrap<T> client = it.next();
						if ( client.deadline < now ) {
							list.add( client );
							it.remove();
						}
					}
				}
				for ( TWrap<T> client : list ) {
					freeClient( client );
				}
			}
		}, MAX_IDLE_MS * 5, MAX_IDLE_MS * 5 );
	}

	@Override
	public Object invoke( Object tproxy, Method method, Object[] args ) throws Throwable
	{

		// first find the thrift methods
		if ( !thriftMethods.contains( method.getName() ) ) {
			throw new IllegalAccessException( "Cannot call this method on a proxied thrift client" );
		}

		TWrap<T> clientWrap = getClient();
		try {
			Throwable cause = null;
			for ( int i = 1; clientWrap != null; i++ ) {
				try {
					return method.invoke( clientWrap.client, args );
				} catch ( InvocationTargetException e ) {
					cause = e.getCause(); // Get original exception
					if ( cause != null && ! ( cause instanceof TTransportException )
							&& ! ( cause instanceof TProtocolException ) ) {
						// If it's not an exception potentially hinting at dead connection, just pass it on
						throw cause;
					}
					// Get rid of broken connection
					freeClient( clientWrap );
					clientWrap = null;
					if ( cause == null ) {
						cause = e;
					}
				}
				// Call the error callback. As long as true is returned, keep retrying
				if ( !errorCallback.thriftError( i, method.getName(), cause ) ) {
					break;
				}
				// Apparently we should retry, get another client
				if ( clientWrap == null ) {
					clientWrap = getClient();
					cause = null;
				}
			}

			// Uh oh
			if ( cause != null )
				throw cause;
			throw new TTransportException( "Could not connect" );
		} finally {
			returnClient( clientWrap );
		}
	}

	private void freeClient( TWrap<T> client )
	{
		try {
			client.client.getInputProtocol().getTransport().close();
		} catch ( Exception e ) {
		}
		try {
			client.client.getOutputProtocol().getTransport().close();
		} catch ( Exception e ) {
		}

	}

	/**
	 * Get an available client connection. Prefer connection from pool,
	 * honoring the max idle time. If none are available, create a fresh
	 * client connection.
	 */
	private TWrap<T> getClient()
	{
		long now = System.currentTimeMillis();
		synchronized ( pool ) {
			TWrap<T> client;
			while ( ( client = pool.poll() ) != null ) {
				if ( client.deadline >= now )
					return client; // Still fresh
				trash.add( client );
			}
		}
		// No usable existing connection, create new client
		LOGGER.debug( "Creating new thrift client" );
		T client = clientFactory.getNewClient();
		if ( client == null )
			return null;
		return new TWrap<T>( client );
	}

	/**
	 * Return a client connection to the pool, updating its last
	 * use timestamp for proper idle timeout handling.
 	 */
	private void returnClient( TWrap<T> client )
	{
		if ( client == null )
			return;
		client.deadline = System.currentTimeMillis() + MAX_IDLE_MS;
		synchronized ( pool ) {
			pool.push( client );
		}
	}
	
	private static class TWrap<T>
	{
		private final T client;
		private long deadline;
		public TWrap(T client)
		{
			this.client = client;
		}
	}

}