001    
002    /*
003     * Copyright (C) 2011 Archie L. Cobbs. All rights reserved.
004     *
005     * $Id: SocketAcceptor.java 239 2012-01-18 21:31:05Z archie.cobbs $
006     */
007    
008    package org.dellroad.stuff.net;
009    
010    import java.io.IOException;
011    import java.net.InetAddress;
012    import java.net.ServerSocket;
013    import java.net.Socket;
014    import java.util.HashSet;
015    
016    import org.slf4j.Logger;
017    import org.slf4j.LoggerFactory;
018    import org.springframework.beans.factory.DisposableBean;
019    import org.springframework.beans.factory.InitializingBean;
020    
021    /**
022     * Spring bean that listens for connections on a TCP socket and spawns a child thread to handle each new connection.
023     * Subclasses must implement {@link #getSocketHandler}. Only the {@link #getPort port} property is required to be set.
024     */
025    public abstract class SocketAcceptor implements InitializingBean, DisposableBean {
026    
027        /**
028         * Default maximum incoming connection queue length ({@value}).
029         *
030         * @see #setBacklog
031         */
032        public static final int DEFAULT_BACKLOG = 50;
033    
034        private static final long NOTIFICATION_INTERVAL = 5 * 1000 * 1000 * 1000L;      // 5 seconds
035    
036        protected final Logger log = LoggerFactory.getLogger(this.getClass());
037    
038        private final HashSet<SocketInfo> connections = new HashSet<SocketInfo>();
039        private InetAddress address;
040        private int backlog = DEFAULT_BACKLOG;
041        private int port;
042        private int maxConnections;
043        private ServerSocket serverSocket;
044        private Thread serverThread;
045        private boolean started;
046    
047        /**
048         * Verifies configuration and invokes {@link #start}.
049         */
050        @Override
051        public void afterPropertiesSet() throws Exception {
052            if (this.port == 0)
053                throw new IllegalArgumentException("TCP port not set");
054            if (this.port < 1 || this.port > 65535)
055                throw new IllegalArgumentException("invalid TCP port " + this.port);
056            if (this.maxConnections < 0)
057                throw new IllegalArgumentException("invalid maxConnections " + this.maxConnections);
058            this.start();
059        }
060    
061        /**
062         * Invokes {@link #stop}.
063         */
064        @Override
065        public void destroy() throws Exception {
066            this.stop();
067        }
068    
069        /**
070         * Get address to listen on.
071         *
072         * @return address to listen on, or null for any
073         */
074        public InetAddress getInetAddress() {
075            return this.address;
076        }
077        public void setInetAddress(InetAddress address) {
078            this.address = address;
079        }
080    
081        /**
082         * Get maximum connect backlog.
083         */
084        public int getBacklog() {
085            return this.backlog;
086        }
087        public void setBacklog(int backlog) {
088            this.backlog = backlog;
089        }
090    
091        /**
092         * Get maximum number of concurrent connections.
093         *
094         * @return max conncurrent connections, or zero for unlimited
095         */
096        public int getMaxConnections() {
097            return this.maxConnections;
098        }
099        public void setMaxConnections(int maxConnections) {
100            this.maxConnections = maxConnections;
101        }
102    
103        /**
104         * Get TCP port to listen on.
105         */
106        public int getPort() {
107            return this.port;
108        }
109        public void setPort(int port) {
110            this.port = port;
111        }
112    
113        /**
114         * Start accepting incoming connections. Does nothing if already started.
115         */
116        public void start() throws IOException {
117            synchronized (this) {
118    
119                // Already started?
120                if (this.started)
121                    return;
122    
123                // Create server thread
124                String addr = this.address != null ? "" + this.address : "*";
125                String threadName = this.getClass().getSimpleName() + "[" + addr + ":" + this.port + "]";
126                this.serverThread = new Thread(threadName) {
127                    @Override
128                    public void run() {
129                        SocketAcceptor.this.run();
130                    }
131                };
132    
133                // Create socket
134                this.serverSocket = this.createServerSocket();
135                if (this.serverSocket == null)
136                    throw new IOException("createServerSocket() returned a null socket");
137            }
138    
139            // Start server thread
140            this.serverThread.start();
141            this.started = true;
142        }
143    
144        private void run() {
145            try {
146                while (true) {
147    
148                    // Block while we've reached our connection limit
149                    synchronized (this) {
150                        boolean logged = false;
151                        while (this.serverSocket != null && this.maxConnections > 0 && this.connections.size() >= this.maxConnections) {
152                            if (!logged) {
153                                this.log.warn(Thread.currentThread().getName() + " has reached connection limit of "
154                                  + this.maxConnections + ", temporarily refusing new connnections");
155                                logged = true;
156                            }
157                            try {
158                                this.wait();
159                            } catch (InterruptedException e) {
160                                // ignore
161                            }
162                        }
163                        if (logged)
164                            this.log.info(Thread.currentThread().getName() + " is accepting new connections again");
165                    }
166    
167                    // Have we been stopped?
168                    ServerSocket serverSocketCopy;
169                    synchronized (this) {
170                        serverSocketCopy = this.serverSocket;
171                    }
172                    if (serverSocketCopy == null)
173                        break;
174    
175                    // Accept a new connection
176                    final Socket socket = serverSocketCopy.accept();
177                    final String socketDesc = socket.getInetAddress().getHostAddress() + ":" + socket.getPort();
178                    this.log.info("accepted new TCP connection from " + socketDesc);
179    
180                    // Get a handler for it
181                    final SocketHandler handler = this.getSocketHandler(socket);
182                    if (handler == null) {
183                        this.log.info("null handler returned by getSocketHandler, closing connection from " + socketDesc);
184                        try {
185                            socket.close();
186                        } catch (IOException e) {
187                            // ignore
188                        }
189                        continue;
190                    }
191                    final SocketInfo socketInfo = new SocketInfo(socket, handler);
192    
193                    // Create a thread that will handle the connection.
194                    Thread handlerThread = new Thread() {
195                        public void run() {
196                            try {
197                                handler.handleConnection(socket);
198                            } catch (Throwable t) {
199                                SocketAcceptor.this.log.error("error handling connection in " + Thread.currentThread().getName(), t);
200                            } finally {
201                                synchronized (SocketAcceptor.this) {
202                                    SocketAcceptor.this.connections.remove(socketInfo);
203                                    SocketAcceptor.this.notifyAll();
204                                }
205                                try {
206                                    socket.close();
207                                } catch (IOException e) {
208                                    // ignore
209                                }
210                            }
211                        }
212                    };
213                    handlerThread.setName(handler.getClass().getSimpleName() + "[" + socketDesc + "]");
214                    socketInfo.setThread(handlerThread);
215    
216                    // Start handler thread and update active connections
217                    synchronized (SocketAcceptor.this) {
218    
219                        // Need to check again for stopped-ness because we released the lock
220                        if (this.serverSocket == null) {
221                            this.log.info("discarding connection from " + socketDesc + " due to shutdown");
222                            try {
223                                socket.close();
224                            } catch (IOException e) {
225                                // ignore
226                            }
227                            break;
228                        }
229    
230                        // Start handler
231                        handlerThread.start();
232                        this.connections.add(socketInfo);
233                    }
234                }
235            } catch (IOException e) {
236                boolean expected;
237                synchronized (this) {
238                    expected = this.serverSocket == null;
239                }
240                if (!expected)
241                    this.log.error("exception in acceptor thread " + Thread.currentThread().getName(), e);
242            } catch (Throwable t) {
243                this.log.error("exception in acceptor thread " + Thread.currentThread().getName(), t);
244            } finally {
245                this.log.info("acceptor thread " + Thread.currentThread().getName() + " exiting");
246                this.closeServerSocket();
247                this.serverThread = null;
248                synchronized (this) {
249                    this.notifyAll();
250                }
251            }
252        }
253    
254        /**
255         * Stop accepting connections. Does nothing if already stopped.
256         *
257         * <p>
258         * Any currently active connections are stopped via {@link SocketHandler#stop SocketHandler.stop()},
259         * and this method waits until all such connections have returned from {@link SocketHandler#handleConnection
260         * SocketHandler.handleConnection()} before returning.
261         */
262        public synchronized void stop() {
263    
264            // Already shut down?
265            if (!this.started)
266                return;
267    
268            // Stop acceptor thread
269            if (this.serverThread != null)
270                this.log.info("stopping acceptor thread");
271            this.closeServerSocket();
272            while (this.serverThread != null) {
273                try {
274                    this.wait();
275                } catch (InterruptedException e) {
276                    // ignore
277                }
278            }
279    
280            // Notify all active connections
281            for (SocketInfo socketInfo : this.connections) {
282    
283                // Notify handler
284                try {
285                    socketInfo.getHandler().stop(socketInfo.getThread(), socketInfo.getSocket());
286                } catch (Throwable t) {
287                    this.log.error("error stopping " + socketInfo.getHandler(), t);
288                }
289    
290                // Close socket
291                try {
292                    socketInfo.getSocket().close();
293                } catch (IOException e) {
294                    // ignore
295                }
296            }
297    
298            // Wait for all active connections to complete
299            long lastTime = System.nanoTime();
300            boolean logged = false;
301            while (!this.connections.isEmpty()) {
302                long nextTime = System.nanoTime();
303                if (!logged || nextTime - lastTime > NOTIFICATION_INTERVAL) {
304                    this.log.info("waiting for " + this.connections.size() + " active connection(s) to stop...");
305                    logged = true;
306                }
307                try {
308                    this.wait();
309                } catch (InterruptedException e) {
310                    // ignore
311                }
312            }
313            if (logged)
314                this.log.info("all active connection(s) have stopped");
315    
316            // Done
317            this.started = false;
318        }
319    
320        /**
321         * Shut down socket.
322         *
323         * @return true if socket was shut down by this invocation, false if it was already shut down
324         */
325        private synchronized void closeServerSocket() {
326            if (this.serverSocket == null)
327                return;
328            try {
329                this.serverSocket.close();
330            } catch (IOException e) {
331                // ignore
332            } finally {
333                this.serverSocket = null;
334            }
335        }
336    
337        /**
338         * Create the server's socket via which connections will be accepted.
339         *
340         * <p>
341         * The implementation in {@link SocketAcceptor} creates the socket and sets the "reuse address" flag.
342         * Subclasses may override.
343         */
344        protected ServerSocket createServerSocket() throws IOException {
345            ServerSocket socket = new ServerSocket(this.port, this.backlog, this.address);
346            socket.setReuseAddress(true);
347            return socket;
348        }
349    
350        /**
351         * Get the {@link SocketHandler} that will handle a new connection using the given socket.
352         *
353         * @return new handler, or <code>null</code> to disconnect the socket immediately
354         */
355        protected abstract SocketHandler getSocketHandler(Socket socket) throws IOException;
356    
357        // Information about one active connection
358        private static class SocketInfo {
359    
360            private final Socket socket;
361            private final SocketHandler handler;
362            private Thread thread;
363    
364            SocketInfo(Socket socket, SocketHandler handler) {
365                this.socket = socket;
366                this.handler = handler;
367            }
368    
369            public Socket getSocket() {
370                return this.socket;
371            }
372    
373            public SocketHandler getHandler() {
374                return this.handler;
375            }
376    
377            public Thread getThread() {
378                return this.thread;
379            }
380            public void setThread(Thread thread) {
381                this.thread = thread;
382            }
383        }
384    }
385