JUnit test rule to reserve a network port

ReservePortsRule.java

/*
 * The MIT License
 *
 * Copyright (c) 2004, The Codehaus
 * Copyright (c) 2015 Mike Chaberski
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy of
 * this software and associated documentation files (the "Software"), to deal in
 * the Software without restriction, including without limitation the rights to
 * use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies
 * of the Software, and to permit persons to whom the Software is furnished to do
 * so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 */

package com.github.mike10004.testhelp;

import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkState;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Lists;
import java.io.IOException;
import java.net.ServerSocket;
import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.TreeMap;
import java.util.UUID;
import java.util.logging.Level;
import java.util.logging.Logger;
import javax.annotation.Nullable;
import org.junit.rules.ExternalResource;

/**
 * Test rule that finds one or more unused ports. Modeled after the Codehaus
 * build-helper-maven-plugin reserve-network-port goal. 
 * 
 * @author Mike Chaberski
 * @author Codehaus
 */
public class ReservePortsRule extends ExternalResource {

    private final ImmutableSet<String> portNames;
    private ImmutableMap<String, Integer> reservedPorts;

    public ReservePortsRule() {
        this(newRandomUniqueString());
    }
    
    public ReservePortsRule(String firstPortName, String...otherPortNames) {
        this(Lists.asList(firstPortName, otherPortNames));
    }
    
    public ReservePortsRule(Collection<String> portNames) {
        this.portNames = ImmutableSet.copyOf(portNames);
        checkArgument(!portNames.isEmpty(), "port name list must be non-empty");
        checkArgument(this.portNames.size() == portNames.size(), "port names collection has duplicate elements");
    }
    
    protected static String newRandomUniqueString() {
        return UUID.randomUUID().toString();
    }
    
    public ImmutableMap<String, Integer> getPorts() {
        checkState(reservedPorts != null, "before() method has not yet been invoked");
        return reservedPorts;
    }
    
    public @Nullable Integer getPort(String portName) {
        checkState(reservedPorts != null, "before() method has not yet been invoked");
        return reservedPorts.get(portName);
    }
    
    public int first() {
        checkState(reservedPorts != null, "before() method has not yet been invoked");
        return reservedPorts.values().iterator().next();
    }
    
    @Override
    protected void before() throws IOException {
        reservedPorts = reservePorts();
    }

    protected ImmutableMap<String, Integer> reservePorts() throws IOException {
        PortReserver portReserver = new PortReserver();
        Map<String, Integer> reservedPorts_ = portReserver.execute(portNames);
        return ImmutableMap.copyOf(reservedPorts_);
    }
    
    @SuppressWarnings("LoggerStringConcat")
    public static class PortReserver
    {
        private static final Integer FIRST_NON_ROOT_PORT_NUMBER = 1024;

        private static final Integer MAX_PORT_NUMBER = 65536;

        private static final Object lock = new Object();

        /**
         * Specify this if you want the port be chosen with a number higher than that one.
         * <p>
         * If {@link #maxPortNumber} is specified, defaults to {@value #FIRST_NON_ROOT_PORT_NUMBER}.
         * </p>
         * 
         * @since 1.8
         */
        private Integer minPortNumber;

        /**
         * Specify this if you want the port be chosen with a number lower than that one.
         * 
         * @since 1.8
         */
        private Integer maxPortNumber;

        protected Logger getLog() {
            return Logger.getLogger(getClass().getName());
        }

        public Map<String, Integer> execute(Iterable<String> portNames) throws IOException
        {
            Map<String, Integer> properties = new TreeMap<>();

            // Reserve the entire block of ports to guarantee we don't get the same port twice
            final List<ServerSocket> sockets = new ArrayList<>();
            final List<Integer> reservedPorts = new ArrayList<>();
            try
            {
                for ( String portName : portNames )
                {
                        final ServerSocket socket = openServerSocket(reservedPorts);
                        sockets.add( socket );

                        final int unusedPort = socket.getLocalPort();
                        properties.put( portName, unusedPort );
                        reservedPorts.add( socket.getLocalPort() );
                        this.getLog().info( "Reserved port " + unusedPort + " for " + portName );
                }

            }
            finally
            {
                // Now free all the ports
                for ( ServerSocket socket : sockets )
                {
                    final int localPort = socket.getLocalPort();
                    try
                    {
                        socket.close();
                    }
                    catch ( IOException e )
                    {
                        this.getLog().log(Level.SEVERE,  "Cannot free reserved port " + localPort, e);
                    }
                }
            }
            return properties;
        }

        private ServerSocket openServerSocket(List<Integer> reservedPorts)
            throws IOException
        {
            if ( minPortNumber == null && maxPortNumber != null )
            {
                getLog().finer( "minPortNumber unspecified: using default value " + FIRST_NON_ROOT_PORT_NUMBER );
                minPortNumber = FIRST_NON_ROOT_PORT_NUMBER;
            }
            if ( minPortNumber != null && maxPortNumber == null )
            {
                getLog().finer( "maxPortNumber unspecified: using default value " + MAX_PORT_NUMBER );
                maxPortNumber = MAX_PORT_NUMBER;
            }
            if ( minPortNumber == null && maxPortNumber == null )
            {
                return new ServerSocket( 0 );
            }
            else
            {
                // Might be synchronizing a bit too largely, but at least that defensive approach should prevent
                // threading issues (essentially possible while put/getting the plugin ctx to get the reserved ports).
                synchronized ( lock )
                {
                    int min = getNextPortNumber(reservedPorts);
                    for ( int port = min;; ++port )
                    {
                        if ( port > maxPortNumber )
                        {
                            throw new IOException( "Unable to find an available port between " + minPortNumber
                                + " and " + maxPortNumber );
                        }
                        try
                        {
                            ServerSocket serverSocket = new ServerSocket( port );
                            return serverSocket;
                        }
                        catch ( IOException ioe )
                        {
                            getLog().log(Level.FINER,  "Tried binding to port " + port + " without success. Trying next port.", ioe );
                        }
                    }
                }
            }
        }

        private int getNextPortNumber(List<Integer> reservedPorts)
        {
            assert minPortNumber != null;

            int nextPort;
            if ( reservedPorts.isEmpty() )
            {
                nextPort = minPortNumber;
            }
            else
            {
                nextPort = findAvailablePortNumber( minPortNumber, reservedPorts );
            }
            reservedPorts.add( nextPort );
            getLog().finer( "Next port: " + nextPort );
            return nextPort;
        }

        /**
         * Returns the first number available, starting at portNumberStartingPoint that's not already in the reservedPorts
         * list.
         * 
         * @param portNumberStartingPoint first port number to start from.
         * @param reservedPorts the ports already reserved.
         * @return first number available not in the given list, starting at the given parameter.
         */
        private int findAvailablePortNumber( Integer portNumberStartingPoint, List<Integer> reservedPorts )
        {
            assert portNumberStartingPoint != null;
            int candidate = portNumberStartingPoint;
            while ( reservedPorts.contains( candidate ) )
            {
                candidate++;
            }
            return candidate;
        }
    }
    
}