Source:
<code>public <Q, R> Stream<QueryResponseMessage<R>> scatterGather(QueryMessage<Q, R> queryMessage, long timeout, TimeUnit timeUnit) { <START> validateShutdown("scatter-gather queries"); <END> QueryMessage<Q, R> interceptedQuery = dispatchInterceptors.intercept(queryMessage); String context = targetContextResolver.resolveContext(interceptedQuery); QueryRequest queryRequest = serializer.serializeRequest(interceptedQuery, SCATTER_GATHER_NUMBER_OF_RESULTS, timeUnit.toMillis(timeout), priorityCalculator.determinePriority(interceptedQuery)); BufferingSpliterator<QueryResponseMessage<R>> resultSpliterator = new BufferingSpliterator<>(Instant.now().plusMillis(timeUnit.toMillis(timeout))); queryService(context) .withDeadlineAfter(timeout, timeUnit) .query(queryRequest, new UpstreamAwareStreamObserver<QueryResponse>() { @Override public void onNext(QueryResponse queryResponse) { logger.debug("Received query response [{}]", queryResponse); if (queryResponse.hasErrorMessage()) { logger.debug("The received query response has error message [{}]", queryResponse.getErrorMessage()); } else { if (!resultSpliterator.put( serializer.deserializeResponse(queryResponse, queryMessage.getResponseType()) )) { getRequestStream().cancel("Cancellation requested by client", null); } } } @Override public void onError(Throwable throwable) { if (!isDeadlineExceeded(throwable)) { logger.info("Received error while waiting for responses", throwable); } resultSpliterator.cancel(throwable); shutdownLatch.decrement(); } @Override public void onCompleted() { resultSpliterator.cancel(null); shutdownLatch.decrement(); } }); shutdownLatch.increment(); return StreamSupport.stream(resultSpliterator, false).onClose(() -> resultSpliterator.cancel(null)); }</code><technical_language>Suggestion: assertRunning assertNotShuttingDown</technical_language>

Target:
public <Q, R> Stream<QueryResponseMessage<R>> scatterGather(QueryMessage<Q, R> queryMessage, long timeout, TimeUnit timeUnit) { shutdownLatch.ifShuttingDown(String.format( "Cannot dispatch new %s as this bus is being shut down", "scatter-gather queries" )); QueryMessage<Q, R> interceptedQuery = dispatchInterceptors.intercept(queryMessage); ShutdownLatch.ActivityHandle queryInTransit = shutdownLatch.registerActivity(); BufferingSpliterator<QueryResponseMessage<R>> resultSpliterator = new BufferingSpliterator<>(Instant.now().plusMillis(timeUnit.toMillis(timeout))); try { String context = targetContextResolver.resolveContext(interceptedQuery); QueryRequest queryRequest = serializer.serializeRequest(interceptedQuery, SCATTER_GATHER_NUMBER_OF_RESULTS, timeUnit.toMillis(timeout), priorityCalculator.determinePriority(interceptedQuery)); QueryServiceGrpc.QueryServiceStub queryService = queryService(context).withDeadlineAfter(timeout, timeUnit); queryService.query(queryRequest, new UpstreamAwareStreamObserver<QueryResponse>() { @Override public void onNext(QueryResponse queryResponse) { logger.debug("Received query response [{}]", queryResponse); if (queryResponse.hasErrorMessage()) { logger.debug("The received query response has error message [{}]", queryResponse.getErrorMessage()); } else { if (!resultSpliterator.put(serializer.deserializeResponse( queryResponse, queryMessage.getResponseType() ))) { getRequestStream().cancel("Cancellation requested by client", null); } } } @Override public void onError(Throwable throwable) { if (!isDeadlineExceeded(throwable)) { logger.info("Received error while waiting for responses", throwable); } resultSpliterator.cancel(throwable); queryInTransit.end(); } @Override public void onCompleted() { resultSpliterator.cancel(null); queryInTransit.end(); } }); } catch (Exception e) { logger.debug("There was a problem issuing a scatter-gather query {}.", interceptedQuery, e); queryInTransit.end(); throw e; } return StreamSupport.stream(resultSpliterator, false).onClose(() -> resultSpliterator.cancel(null)); }