Do not share same socket between two threads at the same time

Question

I have a code in which I am dealing with sockets and I need to make sure that I don't share same socket between two threads. In my below code, I have a background thread which runs every 60 seconds and calls updateLiveSockets() method. In the updateLiveSockets() method, I iterate all the sockets I have and then start pinging them one by one by calling execute method of SendToSocket class and basis on the response I mark them as live or dead.

Now all the reader threads will call getNextSocket() method concurrently to get the next live available socket so it has to be thread safe and I need to make sure all the reader threads should see the same consitent state of SocketHolder and Socket.

Below is my SocketManager class:

public class SocketManager {
  private static final Random random = new Random();
  private final ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor();
  private final Map<Datacenters, List<SocketHolder>> liveSocketsByDatacenter =
      new ConcurrentHashMap<>();
  private final ZContext ctx = new ZContext();

  // ...

  private SocketManager() {
    connectToZMQSockets();
    scheduler.scheduleAtFixedRate(this::updateLiveSockets, 60, 60, TimeUnit.SECONDS);
  }

  // during startup, making a connection and populate once
  private void connectToZMQSockets() {
    Map<Datacenters, List<String>> socketsByDatacenter = Utils.SERVERS;
    for (Map.Entry<Datacenters, List<String>> entry : socketsByDatacenter.entrySet()) {
      List<SocketHolder> addedColoSockets = connect(entry.getValue(), ZMQ.PUSH);
      liveSocketsByDatacenter.put(entry.getKey(), addedColoSockets);
    }
  }

  private List<SocketHolder> connect(List<String> paddes, int socketType) {
    List<SocketHolder> socketList = new ArrayList<>();
    // ....
    return socketList;
  }

  // this method will be called by multiple threads concurrently to get the next live socket
  // is there any concurrency or thread safety issue or race condition here?
  public Optional<SocketHolder> getNextSocket() {
    for (Datacenters dc : Datacenters.getOrderedDatacenters()) {
      Optional<SocketHolder> liveSocket = getLiveSocket(liveSocketsByDatacenter.get(dc));
      if (liveSocket.isPresent()) {
        return liveSocket;
      }
    }
    return Optional.absent();
  }

  private Optional<SocketHolder> getLiveSocket(final List<SocketHolder> listOfEndPoints) {
    if (!CollectionUtils.isEmpty(listOfEndPoints)) {
      // The list of live sockets
      List<SocketHolder> liveOnly = new ArrayList<>(listOfEndPoints.size());
      for (SocketHolder obj : listOfEndPoints) {
        if (obj.isLive()) {
          liveOnly.add(obj);
        }
      }
      if (!liveOnly.isEmpty()) {
        // The list is not empty so we shuffle it an return the first element
        return Optional.of(liveOnly.get(random.nextInt(liveOnly.size()))); // just pick one
      }
    }
    return Optional.absent();
  }

  // runs every 60 seconds to ping all the socket to make sure whether they are alive or not
  private void updateLiveSockets() {
    Map<Datacenters, List<String>> socketsByDatacenter = Utils.SERVERS;

    for (Map.Entry<Datacenters, List<String>> entry : socketsByDatacenter.entrySet()) {
      List<SocketHolder> liveSockets = liveSocketsByDatacenter.get(entry.getKey());
      List<SocketHolder> liveUpdatedSockets = new ArrayList<>();
      for (SocketHolder liveSocket : liveSockets) {
        Socket socket = liveSocket.getSocket();
        String endpoint = liveSocket.getEndpoint();
        Map<byte[], byte[]> holder = populateMap();
        Message message = new Message(holder, Partition.COMMAND);

        // pinging to see whether a socket is live or not
        boolean status = SendToSocket.getInstance().execute(message.getAdd(), holder, socket);
        boolean isLive = (status) ? true : false;

        SocketHolder zmq = new SocketHolder(socket, liveSocket.getContext(), endpoint, isLive);
        liveUpdatedSockets.add(zmq);
      }
      liveSocketsByDatacenter.put(entry.getKey(), Collections.unmodifiableList(liveUpdatedSockets));
    }
  }
}

And here is my SendToSocket class:

// as of now this is only called by "SocketManager" class
public boolean execute(final long padd, final byte[] holder, final Socket socket) {
    boolean sent = sendAsync(padd, encRed, socket, true);
    // if the record was sent successfully, then only sleep for timeout period
    if (sent) {
      try {
        TimeUnit.MILLISECONDS.sleep(500);
      } catch (InterruptedException ex) {
        Thread.currentThread().interrupt();
      }
    }
    return sent;
}

public boolean sendAsync(final long padd, final byte[] encRed, final Socket socket,
  final boolean sockMess) {
    ZMsg msg = new ZMsg();
    msg.add(encRed);
    boolean sent = msg.send(socket);
    msg.destroy();
    retryQueue.put(padd, encRed);
    return sent;
}

// this method will be called by multiple threads concurrently
public boolean sendAsync(final long addr, final byte[] reco) {
  Optional<SocketHolder> liveSockets = SocketManager.getInstance().getNextSocket();
  return sendAsync(addr, reco, liveSockets.get().getSocket(), false);
}

Problem Statement

Now as you can see that I am sharing same socket between two threads. It seems getNextSocket() could return a 0MQ socket to thread A. Concurrently, the timer thread may access the same 0MQ socket to ping it. In this case thread A and the timer thread are mutating the same 0MQ socket, which can lead to problems. So I am trying to find a way so that I can prevent different threads from sending data to the same socket at the same time and mucking up my data.

So I decided to synchronize the socket so that no two threads can access at the same time. Below is the change I made in updateLiveSockets method. I synchronized on the socket in below method:

  // runs every 60 seconds to ping all the socket to make sure whether they are alive or not
  private void updateLiveSockets() {
    Map<Datacenters, List<String>> socketsByDatacenter = Utils.SERVERS;

    for (Map.Entry<Datacenters, List<String>> entry : socketsByDatacenter.entrySet()) {
      List<SocketHolder> liveSockets = liveSocketsByDatacenter.get(entry.getKey());
      List<SocketHolder> liveUpdatedSockets = new ArrayList<>();
      for (SocketHolder liveSocket : liveSockets) {
        Socket socket = liveSocket.getSocket();
        String endpoint = liveSocket.getEndpoint();
        Map<byte[], byte[]> holder = populateMap();
        Message message = new Message(holder, Partition.COMMAND);

        // using the socket as its own lock
        synchronized (socket) {
            // pinging to see whether a socket is live or not
            boolean status = SendToSocket.getInstance().execute(message.getAdd(), holder, socket);
            boolean isLive = (status) ? true : false;

            SocketHolder zmq = new SocketHolder(socket, liveSocket.getContext(), endpoint, isLive);
            liveUpdatedSockets.add(zmq);
        }
      }
      liveSocketsByDatacenter.put(entry.getKey(), Collections.unmodifiableList(liveUpdatedSockets));
    }
  }

And below is the change I made in sendAsync method. In this also I synchronized on socket before sending on it.

public boolean sendAsync(final long padd, final byte[] encRed, final Socket socket,
  final boolean sockMess) {
    ZMsg msg = new ZMsg();
    msg.add(encRed);
    // again, using the socket as its own lock while accessing it
    boolean sent;
    synchronized (socket) {
        sent = msg.send(socket);
    }
    msg.destroy();
    retryQueue.put(padd, encRed);
    return sent;
}

What is the best way by which I can make sure that I am not sharing same sockets between two threads? In general I have around 60 sockets and 20 threads accessing those sockets. I have my above code working fine so opting for a code review to see if there is any better or efficient way to do the same thing? Also is there any corner case or edge case that I have missed which can lead to some bug?

If many threads use the same socket, resources aren't well utilized. Moreover If msg.send(socket); is blocked (technically it shouldn't) all threads waiting for this socket are blocked. So I guess there might be a better way to ensure that every thread uses a different single live socket at the same time instead of synchronization on a particular socket.

Answer 1

Let's start with the "easy" stuff:

• (status) ? true : false is just status (in updateLiveSockets)
• And while we're there: status is only used to assign to isLive. Skip that intermediary and assign directly to isLive
• CollectionUtils.isEmpty(listOfEndPoints) reads much easier as listOfEndpoints.isEmpty(). If that's not semantically the same, isEmpty desperately needs a rename.
• The code you show only uses scheduler in SocketManager's constructor. If you want to ever be able to clean your threads up, you need to expose a method that can end the scheduler's execution. Alternatively you could make the scheduler use "daemon"-threads, which don't prevent the JVM from shutting down.
• The message from sendAsync is put into the retryQueue (which is not defined in the code you show) in all cases. That's most likely wrong.
• The name sendAsync is misleading, because the method is decidedly not asynchronous. Either change the name, or make the method actually async...
• The pluralization in all the code is a mess. If you can't even narrow down the name of the variables into a singular, how do you expect to keep track of singular socket(s) across threads?
• You're overusing Optional in the SocketManager and ignoring it in the calling code. Optional does not provide any safety in your happy-path and it will still blow the code to smithereens in liveSockets.get().getSocket() if the liveSockets is Optional.absent(). Because get will throw an exception and at that point you could've just stuck with standard null and get a NullPointerException....
• SocketManager connects to ZMQ sockets on instantiation. That's bad. I've explained why in an unrelated answer, the gist is: It's a bad idea to perform I/O and other long-running things as well as features of the public API of a class on instantiation, because it's not the expected behaviour of a constructor to do so.

• A lot of the code presented here can be rewritten using Streams Consider getNextSocket():

  public Optional<SocketHolder> getNextSocket() {
      return DataCenters.getOrderedDatacenters()
          .stream()
          .map(dc -> getLiveSocket(liveSocketsByDatacenter.get(dc)))
          .filter(Optional::isPresent)
          .first();
  }
  

  Similar ideas apply for basically all the methods in SocketManager.

After this is fixed, I might be able to follow the "narrative" of your code well enough to actually review it for correctness. As it stands I get too distracted and confused by how it's currently formulated :(