| /* |
| * Copyright (C) 2013 Google Inc. |
| * |
| * Licensed 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 com.google.caliper.runner; |
| |
| import static com.google.common.base.Preconditions.checkState; |
| |
| import com.google.caliper.bridge.OpenedSocket; |
| import com.google.caliper.bridge.StartupAnnounceMessage; |
| import com.google.common.base.Supplier; |
| import com.google.common.collect.Maps; |
| import com.google.common.collect.Multimaps; |
| import com.google.common.collect.SetMultimap; |
| import com.google.common.collect.Sets; |
| import com.google.common.util.concurrent.AbstractExecutionThreadService; |
| import com.google.common.util.concurrent.ListenableFuture; |
| import com.google.common.util.concurrent.Service; |
| import com.google.common.util.concurrent.SettableFuture; |
| |
| import java.io.IOException; |
| import java.net.ServerSocket; |
| import java.net.Socket; |
| import java.net.SocketException; |
| import java.util.Collection; |
| import java.util.Map; |
| import java.util.Set; |
| import java.util.UUID; |
| import java.util.concurrent.locks.Lock; |
| import java.util.concurrent.locks.ReentrantLock; |
| |
| import javax.annotation.concurrent.GuardedBy; |
| import javax.inject.Inject; |
| import javax.inject.Singleton; |
| |
| /** |
| * A {@link Service} that manages a {@link ServerSocket}. |
| * |
| * <p> This service provides two pieces of functionality: |
| * <ol> |
| * <li>It adapts {@link ServerSocket#accept()} to a {@link ListenableFuture} of an opened socket. |
| * <li>It demultiplexes incoming connections based on a {@link StartupAnnounceMessage} that is |
| * sent over the socket. |
| * </ol> |
| * |
| * <p>The {@linkplain State states} of this service are as follows: |
| * <ul> |
| * <li>{@linkplain State#NEW NEW} : Idle state, the {@link ServerSocket} is not open yet. |
| * <li>{@linkplain State#STARTING STARTING} : {@link ServerSocket} is opened |
| * <li>{@linkplain State#RUNNING RUNNING} : We are continuously accepting and parsing connections |
| * from the socket. |
| * <li>{@linkplain State#STOPPING STOPPING} : The server socket is closing and all pending |
| * connection requests are terminated, connection requests will fail immediately. |
| * <li>{@linkplain State#TERMINATED TERMINATED} : Idle state, the socket is closed. |
| * <li>{@linkplain State#FAILED FAILED} : The service will transition to failed if it encounters |
| * any errors while accepting connections or reading from connections. |
| * </ul> |
| * |
| * <p>Note to future self. There have been a few attempts to make it so that it is no longer |
| * necessary to dedicate a thread to this service (basically turn it into an AbstractIdleService). |
| * The general idea has been to make callers to getConnection invoke accept, here is why it didn't |
| * work. |
| * <ul> |
| * <li>If you make getConnection a blocking method that calls accept until it finds the |
| * connection with its id, then there is no way to deal with connections that never arrive. |
| * For example, if the worker crashes before connecting then the thread calling accept will |
| * block forever waiting for it. The only way to unblock a thread stuck on accept() is to |
| * close the socket (this holds for ServerSocketChannels and normal ServerSockets), but we |
| * cannot do that in this case because the socket is a shared resource. |
| * <li>If you make getConnection a non-blocking, polling based method then you expose yourself |
| * to potential deadlocks (due to missed signals) depending on what thread you poll from. |
| * If the polling thread is any of the threads that are involved with processing messages |
| * from the worker I believe there to be a deadlock risk. Basically, if the worker sends |
| * messages over its output streams and then calls Socket.connect, and no printing to stdout |
| * or stderr occurs while connecting. Then if the runner polls, but misses the connection |
| * and then tries to read again, it will deadlock. |
| * </ul> |
| */ |
| @Singleton |
| final class ServerSocketService extends AbstractExecutionThreadService { |
| private enum Source { REQUEST, ACCEPT} |
| |
| private final Lock lock = new ReentrantLock(); |
| |
| /** |
| * Contains futures that have either only been accepted or requested. Once both occur they are |
| * removed from this map. |
| */ |
| @GuardedBy("lock") |
| private final Map<UUID, SettableFuture<OpenedSocket>> halfFinishedConnections = Maps.newHashMap(); |
| |
| /** |
| * Contains the history of connections so we can ensure that each id is only accepted once and |
| * requested once. |
| */ |
| @GuardedBy("lock") |
| private final SetMultimap<Source, UUID> connectionState = Multimaps.newSetMultimap( |
| Maps.<Source, Collection<UUID>>newEnumMap(Source.class), |
| new Supplier<Set<UUID>>(){ |
| @Override public Set<UUID> get() { |
| return Sets.newHashSet(); |
| } |
| }); |
| |
| private ServerSocket serverSocket; |
| |
| @Inject ServerSocketService() {} |
| |
| int getPort() { |
| awaitRunning(); |
| checkState(serverSocket != null, "Socket has not been opened yet"); |
| return serverSocket.getLocalPort(); |
| } |
| |
| /** |
| * Returns a {@link ListenableFuture} for an open connection corresponding to the given id. |
| * |
| * <p>N.B. calling this method 'consumes' the connection and as such calling it twice with the |
| * same id will not work, the second future returned will never complete. Similarly calling it |
| * with an id that does not correspond to a worker trying to connect will also fail. |
| */ |
| public ListenableFuture<OpenedSocket> getConnection(UUID id) { |
| checkState(isRunning(), "You can only get connections from a running service: %s", this); |
| return getConnectionImpl(id, Source.REQUEST); |
| } |
| |
| @Override protected void startUp() throws Exception { |
| serverSocket = new ServerSocket(0 /* bind to any available port */); |
| } |
| |
| @Override protected void run() throws Exception { |
| while (isRunning()) { |
| Socket socket; |
| try { |
| socket = serverSocket.accept(); |
| } catch (SocketException e) { |
| // we were closed |
| return; |
| } |
| OpenedSocket openedSocket = OpenedSocket.fromSocket(socket); |
| |
| UUID id = ((StartupAnnounceMessage) openedSocket.reader().read()).trialId(); |
| // N.B. you should not call set with the lock held, to prevent same thread executors from |
| // running with the lock. |
| getConnectionImpl(id, Source.ACCEPT).set(openedSocket); |
| } |
| } |
| |
| /** |
| * Returns a {@link SettableFuture} from the map of connections. |
| * |
| * <p>This method has the following properties: |
| * <ul> |
| * <li>If the id is present in {@link #connectionState}, this will throw an |
| * {@link IllegalStateException}. |
| * <li>The id and source are recorded in {@link #connectionState} |
| * <li>If the future is already in {@link #halfFinishedConnections}, it is removed and |
| * returned. |
| * <li>If the future is not in {@link #halfFinishedConnections}, a new {@link SettableFuture} |
| * is added and then returned. |
| * |
| * <p>These features together ensure that each connection can only be accepted once, only |
| * requested once and once both have happened it will be removed from |
| * {@link #halfFinishedConnections}. |
| */ |
| private SettableFuture<OpenedSocket> getConnectionImpl(UUID id, Source source) { |
| lock.lock(); |
| try { |
| checkState(connectionState.put(source, id), "Connection for %s has already been %s", |
| id, source); |
| SettableFuture<OpenedSocket> future = halfFinishedConnections.get(id); |
| if (future == null) { |
| future = SettableFuture.create(); |
| halfFinishedConnections.put(id, future); |
| } else { |
| halfFinishedConnections.remove(id); |
| } |
| return future; |
| } finally { |
| lock.unlock(); |
| } |
| } |
| |
| @Override protected void triggerShutdown() { |
| try { |
| serverSocket.close(); |
| } catch (IOException e) { |
| // best effort... |
| } |
| } |
| |
| @Override protected void shutDown() throws Exception { |
| serverSocket.close(); |
| // Now we have either been asked to stop or have failed with some kind of exception, we want to |
| // notify all pending requests, so if there are any references outside of this class they will |
| // notice. |
| lock.lock(); |
| try { |
| for (SettableFuture<OpenedSocket> future : halfFinishedConnections.values()) { |
| future.setException(new Exception("The socket has been closed")); |
| } |
| halfFinishedConnections.clear(); |
| connectionState.clear(); |
| } finally { |
| lock.unlock(); |
| } |
| } |
| } |