Github user jacques-n commented on a diff in the pull request:

    https://github.com/apache/drill/pull/463#discussion_r59433373
  
    --- Diff: 
exec/rpc/src/main/java/org/apache/drill/exec/rpc/RequestIdMap.java ---
    @@ -20,51 +20,82 @@
     import io.netty.buffer.ByteBuf;
     import io.netty.channel.ChannelFuture;
     
    -import java.util.Map;
    -import java.util.concurrent.ConcurrentHashMap;
    +import java.util.concurrent.atomic.AtomicBoolean;
    +import java.util.concurrent.atomic.AtomicInteger;
     
     import org.apache.drill.common.exceptions.UserRemoteException;
     import org.apache.drill.exec.proto.UserBitShared.DrillPBError;
     
    +import com.carrotsearch.hppc.IntObjectHashMap;
    +import com.carrotsearch.hppc.procedures.IntObjectProcedure;
    +import com.google.common.base.Preconditions;
    +
     /**
    - * Manages the creation of rpc futures for a particular socket.
    + * Manages the creation of rpc futures for a particular socket <--> socket
    + * connection. Generally speaking, there will be two threads working with 
this
    + * class (the socket thread and the Request generating thread). 
Synchronization
    + * is simple with the map being the only thing that is protected. 
Everything
    + * else works via Atomic variables.
      */
    -public class CoordinationQueue {
    -  static final org.slf4j.Logger logger = 
org.slf4j.LoggerFactory.getLogger(CoordinationQueue.class);
    +class RequestIdMap {
    +  static final org.slf4j.Logger logger = 
org.slf4j.LoggerFactory.getLogger(RequestIdMap.class);
    +
    +  private final AtomicInteger value = new AtomicInteger();
    +  private final AtomicBoolean acceptMessage = new AtomicBoolean(true);
     
    -  private final PositiveAtomicInteger circularInt = new 
PositiveAtomicInteger();
    -  private final Map<Integer, RpcOutcome<?>> map;
    +  /** Access to map must be protected. **/
    +  private final IntObjectHashMap<RpcOutcome<?>> map;
     
    -  public CoordinationQueue(int segmentSize, int segmentCount) {
    -    map = new ConcurrentHashMap<Integer, RpcOutcome<?>>(segmentSize, 
0.75f, segmentCount);
    +  public RequestIdMap() {
    +    map = new IntObjectHashMap<RpcOutcome<?>>();
       }
     
       void channelClosed(Throwable ex) {
    +    acceptMessage.set(false);
         if (ex != null) {
    -      RpcException e;
    -      if (ex instanceof RpcException) {
    -        e = (RpcException) ex;
    -      } else {
    -        e = new RpcException(ex);
    +      final RpcException e = RpcException.mapException(ex);
    +      synchronized (map) {
    +        map.forEach(new Closer(e));
    +        map.clear();
           }
    -      for (RpcOutcome<?> f : map.values()) {
    -        f.setException(e);
    +    }
    +  }
    +
    +  private class Closer implements IntObjectProcedure<RpcOutcome<?>> {
    +    final RpcException exception;
    +
    +    public Closer(RpcException exception) {
    +      this.exception = exception;
    +    }
    +
    +    @Override
    +    public void apply(int key, RpcOutcome<?> value) {
    +      try{
    +        value.setException(exception);
    +      }catch(Exception e){
    +        logger.warn("Failure while attempting to fail rpc response.", e);
           }
         }
    +
       }
     
    -  public <V> ChannelListenerWithCoordinationId get(RpcOutcomeListener<V> 
handler, Class<V> clazz, RemoteConnection connection) {
    -    int i = circularInt.getNext();
    +  public <V> ChannelListenerWithCoordinationId 
createNewRpcListener(RpcOutcomeListener<V> handler, Class<V> clazz,
    +      RemoteConnection connection) {
    +    int i = value.incrementAndGet();
         RpcListener<V> future = new RpcListener<V>(handler, clazz, i, 
connection);
    -    Object old = map.put(i, future);
    -    if (old != null) {
    -      throw new IllegalStateException(
    -          "You attempted to reuse a coordination id when the previous 
coordination id has not been removed.  This is likely rpc future callback 
memory leak.");
    +    final Object old;
    +    synchronized (map) {
    +      Preconditions.checkArgument(acceptMessage.get(),
    +          "Attempted to send a message when connection is no longer 
valid.");
    +      old = map.put(i, future);
         }
    +    Preconditions.checkArgument(old == null,
    --- End diff --
    
    This is an assertion to ensure that there isn't a bug some place.  


---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at [email protected] or file a JIRA ticket
with INFRA.
---

Reply via email to