88
99package org .elasticsearch .tasks ;
1010
11- import com .carrotsearch .hppc .ObjectIntHashMap ;
12- import com .carrotsearch .hppc .ObjectIntMap ;
13-
1411import org .apache .logging .log4j .LogManager ;
1512import org .apache .logging .log4j .Logger ;
1613import org .apache .logging .log4j .message .ParameterizedMessage ;
6057import java .util .concurrent .atomic .AtomicBoolean ;
6158import java .util .concurrent .atomic .AtomicLong ;
6259import java .util .function .Consumer ;
63- import java .util .stream .Collectors ;
64- import java .util .stream .StreamSupport ;
6560
6661import static org .elasticsearch .core .TimeValue .timeValueMillis ;
6762import static org .elasticsearch .http .HttpTransportSettings .SETTING_HTTP_MAX_HEADER_SIZE ;
@@ -506,7 +501,7 @@ private static class CancellableTaskHolder {
506501 private final CancellableTask task ;
507502 private boolean finished = false ;
508503 private List <Runnable > cancellationListeners = null ;
509- private ObjectIntMap <Transport .Connection > childTasksPerConnection = null ;
504+ private Map <Transport .Connection , Integer > childTasksPerConnection = null ;
510505 private String banChildrenReason ;
511506 private List <Runnable > childTaskCompletedListeners = null ;
512507
@@ -587,15 +582,15 @@ synchronized void registerChildConnection(Transport.Connection connection) {
587582 throw new TaskCancelledException ("parent task was cancelled [" + banChildrenReason + ']' );
588583 }
589584 if (childTasksPerConnection == null ) {
590- childTasksPerConnection = new ObjectIntHashMap <>();
585+ childTasksPerConnection = new HashMap <>();
591586 }
592- childTasksPerConnection .addTo (connection , 1 );
587+ childTasksPerConnection .merge (connection , 1 , Integer :: sum );
593588 }
594589
595590 void unregisterChildConnection (Transport .Connection node ) {
596591 final List <Runnable > listeners ;
597592 synchronized (this ) {
598- if (childTasksPerConnection .addTo (node , -1 ) == 0 ) {
593+ if (childTasksPerConnection .merge (node , -1 , Integer :: sum ) == 0 ) {
599594 childTasksPerConnection .remove (node );
600595 }
601596 if (childTasksPerConnection .isEmpty () == false || this .childTaskCompletedListeners == null ) {
@@ -617,9 +612,7 @@ Set<Transport.Connection> startBan(String reason, Runnable onChildTasksCompleted
617612 if (childTasksPerConnection == null ) {
618613 pendingChildConnections = Collections .emptySet ();
619614 } else {
620- pendingChildConnections = StreamSupport .stream (childTasksPerConnection .spliterator (), false )
621- .map (e -> e .key )
622- .collect (Collectors .toUnmodifiableSet ());
615+ pendingChildConnections = Set .copyOf (childTasksPerConnection .keySet ());
623616 }
624617 if (pendingChildConnections .isEmpty ()) {
625618 assert childTaskCompletedListeners == null ;
0 commit comments