content: /*
   * Licensed to the Apache Software Foundation (ASF) under one or more
   * contributor license agreements. See the NOTICE file distributed with
   * this work for additional information regarding copyright ownership.
   * The ASF licenses this file to You 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 org.apache.kafka.common.test;
  
  import org.apache.kafka.common.utils.Utils;
  import org.apache.kafka.server.ServerSocketFactory;
  
  import java.io.IOException;
  import java.net.InetSocketAddress;
  import java.nio.channels.ServerSocketChannel;
  import java.util.HashMap;
  import java.util.HashSet;
  import java.util.Map;
  import java.util.Map.Entry;
  import java.util.Set;
  
  public class PreboundSocketFactoryManager implements AutoCloseable {
  
      private class PreboundSocketFactory implements ServerSocketFactory {
          private final int nodeId;
  
          private PreboundSocketFactory(int nodeId) {
              this.nodeId = nodeId;
          }
  
          @Override
          public ServerSocketChannel openServerSocket(
                  String listenerName,
                  InetSocketAddress socketAddress,
                  int listenBacklogSize,
                  int recvBufferSize
          ) throws IOException {
              ServerSocketChannel socketChannel = getSocketForListenerAndMarkAsUsed(
                  nodeId,
                  listenerName);
  
              if (socketChannel != null) {
                  if (socketChannel.isOpen()) {
                      return socketChannel;
                  }
                  // When restarting components(e.g. controllers, brokers) in tests, we want to reuse the same
                  // port that was previously allocated to maintain consistent addressing
                  // so the client can reconnect to the same port.
                  // Since those components would close the socket when they are restarted,
                  // we need to rebind the socket to the same port.
                  socketAddress = new InetSocketAddress(socketAddress.getHostString(), socketChannel.socket().getLocalPort());
                  socketChannel = ServerSocketFactory.INSTANCE.openServerSocket(
                          listenerName,
                          socketAddress,
                          listenBacklogSize,
                          recvBufferSize);
                  return socketChannel;
              }
              return ServerSocketFactory.INSTANCE.openServerSocket(
                  listenerName,
                  socketAddress,
                  listenBacklogSize,
                  recvBufferSize);
          }
      }
  
      /**
       * True if this manager is closed.
       */
      private boolean closed = false;
  
      /**
       * Maps node IDs to socket factory objects.
       * Protected by the object lock.
       */
      private final Map<Integer, PreboundSocketFactory> factories = new HashMap<>();
  
      /**
       * Maps node IDs to maps of listener names to ports.
       * Protected by the object lock.
       */
      private final Map<Integer, Map<String, ServerSocketChannel>> sockets = new HashMap<>();
  
      /**
       * Maps node IDs to set of the listeners that were used.
       * Protected by the object lock.
       */
      private final Map<Integer, Set<String>> usedSockets = new HashMap<>();
  
      /**
       * Get a socket from this manager, mark it as used, and return it.
       *
       * @param nodeId        The ID of the node.
       * @param listener      The listener for the socket.
       *
       * @return              null if the socket was not found; the socket, otherwise.
       */
      public synchronized ServerSocketChannel getSocketForListenerAndMarkAsUsed(
          int nodeId,
          String listener
      ) {
          Map<String, ServerSocketChannel> socketsForNode = sockets.get(nodeId);
          if (socketsForNode == null) {
              return null;
          }
          ServerSocketChannel socket = socketsForNode.get(listener);
          if (socket == null) {
              return null;
          }
          usedSockets.computeIfAbsent(nodeId, __ -> new HashSet<>()).add(listener);
          return socket;
      }
  
      /**
       * Get or create a socket factory object associated with a given node ID.
       *
       * @param nodeId        The ID of the node.
       *
       * @return              The socket factory.
       */
      public synchronized ServerSocketFactory getOrCreateSocketFactory(int nodeId) {
          return factories.computeIfAbsent(nodeId, __ -> new PreboundSocketFactory(nodeId));
      }
  
      /**
       * Get a specific port number. The port will be created if it does not already exist.
       *
       * @param nodeId        The ID of the node.
       * @param listener      The listener for the socket.
       *
       * @return              The port number.
       */
      public synchronized int getOrCreatePortForListener(
          int nodeId,
          String listener
      ) throws IOException {
          Map<String, ServerSocketChannel> socketsForNode =
              sockets.computeIfAbsent(nodeId, __ -> new HashMap<>());
          ServerSocketChannel socketChannel = socketsForNode.get(listener);
          if (socketChannel == null) {
              if (closed) {
                  throw new RuntimeException("Cannot open new socket: manager is closed.");
              }
              socketChannel = ServerSocketFactory.INSTANCE.openServerSocket(
                  listener,
                  new InetSocketAddress(0),
                  -1,
                  -1);
              socketsForNode.put(listener, socketChannel);
          }
          InetSocketAddress socketAddress = (InetSocketAddress) socketChannel.getLocalAddress();
          return socketAddress.getPort();
      }
  
      @Override
      public synchronized void close() throws Exception {
          if (closed) {
              return;
          }
          closed = true;
          // Close all sockets that haven't been used by a SocketServer. (We don't want to close the
          // ones that have been used by a SocketServer because that is the responsibility of that
          // SocketServer.)
          for (Entry<Integer, Map<String, ServerSocketChannel>> socketsEntry : sockets.entrySet()) {
              Set<String> usedListeners = usedSockets.getOrDefault(
                  socketsEntry.getKey(), Set.of());
              for (Entry<String, ServerSocketChannel> entry : socketsEntry.getValue().entrySet()) {
                  if (!usedListeners.contains(entry.getKey())) {
                      Utils.closeQuietly(entry.getValue(), "serverSocketChannel");
                  }
              }
          }
      }
  }
  