gobblin.util.PortUtils.java Source code

Introduction

Here is the source code for gobblin.util.PortUtils.java

Source

/*
 * Copyright (C) 2014-2016 LinkedIn Corp. All rights reserved.
 *
 * 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.
 */

package gobblin.util;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Optional;
import com.google.common.base.Preconditions;
import com.google.common.collect.BiMap;
import com.google.common.collect.HashBiMap;
import com.google.common.collect.Maps;

import java.net.ServerSocket;
import java.util.Map;
import java.util.concurrent.ConcurrentMap;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

public class PortUtils {
    public static final int MINIMUM_PORT = 1025;
    public static final int MAXIMUM_PORT = 65535;
    private static final Pattern PORT_REGEX = Pattern
            .compile("\\$\\{PORT_(?>(?>\\?(\\d+))|(?>(\\d+)\\?)|(\\d+|\\?))\\}");
    private final PortLocator portLocator;
    private final ConcurrentMap<Integer, Boolean> assignedPorts;

    public PortUtils() {
        this(new ServerSocketPortLocator());
    }

    @VisibleForTesting
    PortUtils(PortLocator locator) {
        this.portLocator = locator;
        this.assignedPorts = Maps.newConcurrentMap();
    }

    /**
     * Replaces any port tokens in the specified string.
     *
     * NOTE: Tokens can be in the following forms:
     *   1. ${PORT_123}
     *   2. ${PORT_?123}
     *   3. ${PORT_123?}
     *   4. ${PORT_?}
     *
     * @param value The string in which to replace port tokens.
     * @return The replaced string.
     */
    public String replacePortTokens(String value) {
        BiMap<String, Optional<Integer>> portMappings = HashBiMap.create();
        Matcher regexMatcher = PORT_REGEX.matcher(value);
        while (regexMatcher.find()) {
            String token = regexMatcher.group(0);
            if (!portMappings.containsKey(token)) {
                Optional<Integer> portStart = Optional.absent();
                Optional<Integer> portEnd = Optional.absent();
                String unboundedStart = regexMatcher.group(1);
                if (unboundedStart != null) {
                    int requestedEndPort = Integer.parseInt(unboundedStart);
                    Preconditions.checkArgument(requestedEndPort <= PortUtils.MAXIMUM_PORT);
                    portEnd = Optional.of(requestedEndPort);
                } else {
                    String unboundedEnd = regexMatcher.group(2);
                    if (unboundedEnd != null) {
                        int requestedStartPort = Integer.parseInt(unboundedEnd);
                        Preconditions.checkArgument(requestedStartPort >= PortUtils.MINIMUM_PORT);
                        portStart = Optional.of(requestedStartPort);
                    } else {
                        String absolute = regexMatcher.group(3);
                        if (!"?".equals(absolute)) {
                            int requestedPort = Integer.parseInt(absolute);
                            Preconditions.checkArgument(requestedPort >= PortUtils.MINIMUM_PORT
                                    && requestedPort <= PortUtils.MAXIMUM_PORT);
                            portStart = Optional.of(requestedPort);
                            portEnd = Optional.of(requestedPort);
                        }
                    }
                }
                Optional<Integer> port = takePort(portStart, portEnd);
                portMappings.put(token, port);
            }
        }
        for (Map.Entry<String, Optional<Integer>> port : portMappings.entrySet()) {
            if (port.getValue().isPresent()) {
                value = value.replace(port.getKey(), port.getValue().get().toString());
            }
        }
        return value;
    }

    /**
     * Finds an open port. {@param portStart} and {@param portEnd} can be absent
     *
     * ______________________________________________________
     * | portStart | portEnd  | takenPort                   |
     * |-----------|----------|-----------------------------|
     * |  absent   | absent   | random                      |
     * |  absent   | provided | 1024 < port <= portEnd      |
     * |  provided | absent   | portStart <= port <= 65535  |
     * |  provided | provided | portStart = port = portEnd  |
     * ------------------------------------------------------
     *
     * @param portStart the inclusive starting port
     * @param portEnd the inclusive ending port
     * @return The selected open port.
     */
    private synchronized Optional<Integer> takePort(Optional<Integer> portStart, Optional<Integer> portEnd) {
        if (!portStart.isPresent() && !portEnd.isPresent()) {
            for (int i = 0; i < 65535; i++) {
                try {
                    int port = this.portLocator.random();
                    Boolean wasAssigned = assignedPorts.putIfAbsent(port, true);
                    if (wasAssigned == null || !wasAssigned) {
                        return Optional.of(port);
                    }
                } catch (Exception ignored) {
                }
            }
        }

        for (int port = portStart.or(MINIMUM_PORT); port <= portEnd.or(MAXIMUM_PORT); port++) {
            try {
                this.portLocator.specific(port);
                Boolean wasAssigned = assignedPorts.putIfAbsent(port, true);
                if (wasAssigned == null || !wasAssigned) {
                    return Optional.of(port);
                }
            } catch (Exception ignored) {
            }
        }

        throw new RuntimeException(String.format("No open port could be found for %s to %s", portStart, portEnd));
    }

    @VisibleForTesting
    interface PortLocator {
        int random() throws Exception;

        int specific(int port) throws Exception;
    }

    private static class ServerSocketPortLocator implements PortLocator {
        @Override
        public int random() throws Exception {
            try (ServerSocket serverSocket = new ServerSocket(0)) {
                return serverSocket.getLocalPort();
            }
        }

        @Override
        public int specific(int port) throws Exception {
            try (ServerSocket serverSocket = new ServerSocket(port)) {
                return serverSocket.getLocalPort();
            }
        }
    }
}