@Test public void unsupportedQueryInstruction() { FakeStreamObserver<QueryProviderInbound> responseStream = new FakeStreamObserver<>(); StreamObserver<QueryProviderOutbound> requestStream = testSubject.openStream(responseStream); String instructionId = "instructionId"; requestStream.onNext(QueryProviderOutbound.newBuilder() .setInstructionId(instructionId) .build()); InstructionAck ack = responseStream.values().get(responseStream.values().size() - 1) .getAck(); assertEquals(instructionId, ack.getInstructionId()); assertTrue(ack.hasError()); assertEquals(ErrorCode.UNSUPPORTED_INSTRUCTION.getCode(), ack.getError().getErrorCode()); }
@Override public StreamObserver<QueryProviderOutbound> openStream( StreamObserver<QueryProviderInbound> inboundStreamObserver) { String context = contextProvider.getContext(); SendingStreamObserver<QueryProviderInbound> wrappedQueryProviderInboundObserver = new SendingStreamObserver<>( inboundStreamObserver); return new ReceivingStreamObserver<QueryProviderOutbound>(logger) { private final AtomicReference<String> clientIdRef = new AtomicReference<>(); private final AtomicReference<GrpcQueryDispatcherListener> listener = new AtomicReference<>(); private final AtomicReference<ClientStreamIdentification> clientRef = new AtomicReference<>(); private final AtomicReference<QueryHandler<QueryProviderInbound>> queryHandler = new AtomicReference<>(); @Override protected void consume(QueryProviderOutbound queryProviderOutbound) { switch (queryProviderOutbound.getRequestCase()) { case SUBSCRIBE: instructionAckSource.sendSuccessfulAck(queryProviderOutbound.getInstructionId(), wrappedQueryProviderInboundObserver); QuerySubscription subscription = queryProviderOutbound.getSubscribe(); checkInitClient(subscription.getClientId(), subscription.getComponentName()); String clientStreamId = clientRef.get().getClientStreamId(); logger.debug("{}: Subscribe Query {} for {}", context, subscription.getQuery(), subscription.getClientId()); SubscribeQuery subscribeQuery = new SubscribeQuery(context, clientStreamId, subscription, queryHandler.get()); eventPublisher.publishEvent( subscribeQuery); break; case UNSUBSCRIBE: instructionAckSource.sendSuccessfulAck(queryProviderOutbound.getInstructionId(), wrappedQueryProviderInboundObserver); if (clientRef.get() != null) { QuerySubscription unsubscribe = queryProviderOutbound.getUnsubscribe(); UnsubscribeQuery unsubscribeQuery = new UnsubscribeQuery(context, clientRef.get() .getClientStreamId(), unsubscribe, false); eventPublisher.publishEvent(unsubscribeQuery); } break; case FLOW_CONTROL: instructionAckSource.sendSuccessfulAck(queryProviderOutbound.getInstructionId(), wrappedQueryProviderInboundObserver); flowControl(queryProviderOutbound.getFlowControl()); break; case QUERY_RESPONSE: instructionAckSource.sendSuccessfulAck(queryProviderOutbound.getInstructionId(), wrappedQueryProviderInboundObserver); queryDispatcher.handleResponse(queryProviderOutbound.getQueryResponse(), clientRef.get().getClientStreamId(), clientIdRef.get(), false); break; case QUERY_COMPLETE: instructionAckSource.sendSuccessfulAck(queryProviderOutbound.getInstructionId(), wrappedQueryProviderInboundObserver); queryDispatcher.handleComplete(queryProviderOutbound.getQueryComplete().getRequestId(), clientRef.get().getClientStreamId(), clientIdRef.get(), false); break; case SUBSCRIPTION_QUERY_RESPONSE: instructionAckSource.sendSuccessfulAck(queryProviderOutbound.getInstructionId(), wrappedQueryProviderInboundObserver); SubscriptionQueryResponse response = queryProviderOutbound.getSubscriptionQueryResponse(); eventPublisher.publishEvent(new SubscriptionQueryResponseReceived(response, () -> wrappedQueryProviderInboundObserver .onNext(unsubscribeMessage(response.getSubscriptionIdentifier())))); break; case ACK: InstructionAck ack = queryProviderOutbound.getAck(); if (isUnsupportedInstructionErrorResult(ack)) { logger.warn("Unsupported instruction sent to the client {} of context {}.", clientRef.get().getClientStreamId(), context); } else { logger.trace("Received instruction ack from the client {} of context {}. Result {}.", clientRef.get().getClientStreamId(), context, ack); } break; default: instructionAckSource.sendUnsupportedInstruction(queryProviderOutbound.getInstructionId(), topology.getMe().getName(), wrappedQueryProviderInboundObserver); break; } } private void flowControl(FlowControl flowControl) { initClientReference(flowControl.getClientId()); if (listener.compareAndSet(null, new GrpcQueryDispatcherListener(queryDispatcher, clientRef.get().toString(), wrappedQueryProviderInboundObserver, processingThreads))) { dispatcherListeners.put(clientRef.get(), listener.get()); } listener.get().addPermits(flowControl.getPermits()); } private void initClientReference(String clientId) { String clientStreamId = clientId + "." + UUID.randomUUID().toString(); if (clientRef.compareAndSet(null, new ClientStreamIdentification(context, clientStreamId))) { clientIdRegistry.register(clientStreamId, clientId, ClientIdRegistry.ConnectionType.QUERY); } clientIdRef.compareAndSet(null, clientId); } private void checkInitClient(String clientId, String componentName) { initClientReference(clientId); queryHandler.compareAndSet(null, new DirectQueryHandler(wrappedQueryProviderInboundObserver, clientRef.get(), componentName, clientId)); } @Override protected String sender() { return clientRef.toString(); } @Override public void onError(Throwable cause) { if (!ExceptionUtils.isCancelled(cause)) { logger.warn("{}: Error on connection from subscriber - {}", clientRef, cause.getMessage()); } cleanup(); } private void cleanup() { if (clientRef.get() != null) { String clientStreamId = clientRef.get().getClientStreamId(); String clientId = this.clientIdRef.get(); clientIdRegistry.unregister(clientStreamId, ClientIdRegistry.ConnectionType.QUERY); eventPublisher.publishEvent(new QueryHandlerDisconnected(context, clientId, clientStreamId)); } if (listener.get() != null) { dispatcherListeners.remove(clientRef.get()); listener.get().cancel(); } StreamObserverUtils.complete(inboundStreamObserver); } @Override public void onCompleted() { cleanup(); } }; }
QueryService extends QueryServiceGrpc.QueryServiceImplBase implements AxonServerClientService { @Override public StreamObserver<QueryProviderOutbound> openStream( StreamObserver<QueryProviderInbound> inboundStreamObserver) { String context = contextProvider.getContext(); SendingStreamObserver<QueryProviderInbound> wrappedQueryProviderInboundObserver = new SendingStreamObserver<>( inboundStreamObserver); return new ReceivingStreamObserver<QueryProviderOutbound>(logger) { private final AtomicReference<String> clientIdRef = new AtomicReference<>(); private final AtomicReference<GrpcQueryDispatcherListener> listener = new AtomicReference<>(); private final AtomicReference<ClientStreamIdentification> clientRef = new AtomicReference<>(); private final AtomicReference<QueryHandler<QueryProviderInbound>> queryHandler = new AtomicReference<>(); @Override protected void consume(QueryProviderOutbound queryProviderOutbound) { switch (queryProviderOutbound.getRequestCase()) { case SUBSCRIBE: instructionAckSource.sendSuccessfulAck(queryProviderOutbound.getInstructionId(), wrappedQueryProviderInboundObserver); QuerySubscription subscription = queryProviderOutbound.getSubscribe(); checkInitClient(subscription.getClientId(), subscription.getComponentName()); String clientStreamId = clientRef.get().getClientStreamId(); logger.debug("{}: Subscribe Query {} for {}", context, subscription.getQuery(), subscription.getClientId()); SubscribeQuery subscribeQuery = new SubscribeQuery(context, clientStreamId, subscription, queryHandler.get()); eventPublisher.publishEvent( subscribeQuery); break; case UNSUBSCRIBE: instructionAckSource.sendSuccessfulAck(queryProviderOutbound.getInstructionId(), wrappedQueryProviderInboundObserver); if (clientRef.get() != null) { QuerySubscription unsubscribe = queryProviderOutbound.getUnsubscribe(); UnsubscribeQuery unsubscribeQuery = new UnsubscribeQuery(context, clientRef.get() .getClientStreamId(), unsubscribe, false); eventPublisher.publishEvent(unsubscribeQuery); } break; case FLOW_CONTROL: instructionAckSource.sendSuccessfulAck(queryProviderOutbound.getInstructionId(), wrappedQueryProviderInboundObserver); flowControl(queryProviderOutbound.getFlowControl()); break; case QUERY_RESPONSE: instructionAckSource.sendSuccessfulAck(queryProviderOutbound.getInstructionId(), wrappedQueryProviderInboundObserver); queryDispatcher.handleResponse(queryProviderOutbound.getQueryResponse(), clientRef.get().getClientStreamId(), clientIdRef.get(), false); break; case QUERY_COMPLETE: instructionAckSource.sendSuccessfulAck(queryProviderOutbound.getInstructionId(), wrappedQueryProviderInboundObserver); queryDispatcher.handleComplete(queryProviderOutbound.getQueryComplete().getRequestId(), clientRef.get().getClientStreamId(), clientIdRef.get(), false); break; case SUBSCRIPTION_QUERY_RESPONSE: instructionAckSource.sendSuccessfulAck(queryProviderOutbound.getInstructionId(), wrappedQueryProviderInboundObserver); SubscriptionQueryResponse response = queryProviderOutbound.getSubscriptionQueryResponse(); eventPublisher.publishEvent(new SubscriptionQueryResponseReceived(response, () -> wrappedQueryProviderInboundObserver .onNext(unsubscribeMessage(response.getSubscriptionIdentifier())))); break; case ACK: InstructionAck ack = queryProviderOutbound.getAck(); if (isUnsupportedInstructionErrorResult(ack)) { logger.warn("Unsupported instruction sent to the client {} of context {}.", clientRef.get().getClientStreamId(), context); } else { logger.trace("Received instruction ack from the client {} of context {}. Result {}.", clientRef.get().getClientStreamId(), context, ack); } break; default: instructionAckSource.sendUnsupportedInstruction(queryProviderOutbound.getInstructionId(), topology.getMe().getName(), wrappedQueryProviderInboundObserver); break; } } private void flowControl(FlowControl flowControl) { initClientReference(flowControl.getClientId()); if (listener.compareAndSet(null, new GrpcQueryDispatcherListener(queryDispatcher, clientRef.get().toString(), wrappedQueryProviderInboundObserver, processingThreads))) { dispatcherListeners.put(clientRef.get(), listener.get()); } listener.get().addPermits(flowControl.getPermits()); } private void initClientReference(String clientId) { String clientStreamId = clientId + "." + UUID.randomUUID().toString(); if (clientRef.compareAndSet(null, new ClientStreamIdentification(context, clientStreamId))) { clientIdRegistry.register(clientStreamId, clientId, ClientIdRegistry.ConnectionType.QUERY); } clientIdRef.compareAndSet(null, clientId); } private void checkInitClient(String clientId, String componentName) { initClientReference(clientId); queryHandler.compareAndSet(null, new DirectQueryHandler(wrappedQueryProviderInboundObserver, clientRef.get(), componentName, clientId)); } @Override protected String sender() { return clientRef.toString(); } @Override public void onError(Throwable cause) { if (!ExceptionUtils.isCancelled(cause)) { logger.warn("{}: Error on connection from subscriber - {}", clientRef, cause.getMessage()); } cleanup(); } private void cleanup() { if (clientRef.get() != null) { String clientStreamId = clientRef.get().getClientStreamId(); String clientId = this.clientIdRef.get(); clientIdRegistry.unregister(clientStreamId, ClientIdRegistry.ConnectionType.QUERY); eventPublisher.publishEvent(new QueryHandlerDisconnected(context, clientId, clientStreamId)); } if (listener.get() != null) { dispatcherListeners.remove(clientRef.get()); listener.get().cancel(); } StreamObserverUtils.complete(inboundStreamObserver); } @Override public void onCompleted() { cleanup(); } }; } }
QueryService extends QueryServiceGrpc.QueryServiceImplBase implements AxonServerClientService { @Override public StreamObserver<QueryProviderOutbound> openStream( StreamObserver<QueryProviderInbound> inboundStreamObserver) { String context = contextProvider.getContext(); SendingStreamObserver<QueryProviderInbound> wrappedQueryProviderInboundObserver = new SendingStreamObserver<>( inboundStreamObserver); return new ReceivingStreamObserver<QueryProviderOutbound>(logger) { private final AtomicReference<String> clientIdRef = new AtomicReference<>(); private final AtomicReference<GrpcQueryDispatcherListener> listener = new AtomicReference<>(); private final AtomicReference<ClientStreamIdentification> clientRef = new AtomicReference<>(); private final AtomicReference<QueryHandler<QueryProviderInbound>> queryHandler = new AtomicReference<>(); @Override protected void consume(QueryProviderOutbound queryProviderOutbound) { switch (queryProviderOutbound.getRequestCase()) { case SUBSCRIBE: instructionAckSource.sendSuccessfulAck(queryProviderOutbound.getInstructionId(), wrappedQueryProviderInboundObserver); QuerySubscription subscription = queryProviderOutbound.getSubscribe(); checkInitClient(subscription.getClientId(), subscription.getComponentName()); String clientStreamId = clientRef.get().getClientStreamId(); logger.debug("{}: Subscribe Query {} for {}", context, subscription.getQuery(), subscription.getClientId()); SubscribeQuery subscribeQuery = new SubscribeQuery(context, clientStreamId, subscription, queryHandler.get()); eventPublisher.publishEvent( subscribeQuery); break; case UNSUBSCRIBE: instructionAckSource.sendSuccessfulAck(queryProviderOutbound.getInstructionId(), wrappedQueryProviderInboundObserver); if (clientRef.get() != null) { QuerySubscription unsubscribe = queryProviderOutbound.getUnsubscribe(); UnsubscribeQuery unsubscribeQuery = new UnsubscribeQuery(context, clientRef.get() .getClientStreamId(), unsubscribe, false); eventPublisher.publishEvent(unsubscribeQuery); } break; case FLOW_CONTROL: instructionAckSource.sendSuccessfulAck(queryProviderOutbound.getInstructionId(), wrappedQueryProviderInboundObserver); flowControl(queryProviderOutbound.getFlowControl()); break; case QUERY_RESPONSE: instructionAckSource.sendSuccessfulAck(queryProviderOutbound.getInstructionId(), wrappedQueryProviderInboundObserver); queryDispatcher.handleResponse(queryProviderOutbound.getQueryResponse(), clientRef.get().getClientStreamId(), clientIdRef.get(), false); break; case QUERY_COMPLETE: instructionAckSource.sendSuccessfulAck(queryProviderOutbound.getInstructionId(), wrappedQueryProviderInboundObserver); queryDispatcher.handleComplete(queryProviderOutbound.getQueryComplete().getRequestId(), clientRef.get().getClientStreamId(), clientIdRef.get(), false); break; case SUBSCRIPTION_QUERY_RESPONSE: instructionAckSource.sendSuccessfulAck(queryProviderOutbound.getInstructionId(), wrappedQueryProviderInboundObserver); SubscriptionQueryResponse response = queryProviderOutbound.getSubscriptionQueryResponse(); eventPublisher.publishEvent(new SubscriptionQueryResponseReceived(response, () -> wrappedQueryProviderInboundObserver .onNext(unsubscribeMessage(response.getSubscriptionIdentifier())))); break; case ACK: InstructionAck ack = queryProviderOutbound.getAck(); if (isUnsupportedInstructionErrorResult(ack)) { logger.warn("Unsupported instruction sent to the client {} of context {}.", clientRef.get().getClientStreamId(), context); } else { logger.trace("Received instruction ack from the client {} of context {}. Result {}.", clientRef.get().getClientStreamId(), context, ack); } break; default: instructionAckSource.sendUnsupportedInstruction(queryProviderOutbound.getInstructionId(), topology.getMe().getName(), wrappedQueryProviderInboundObserver); break; } } private void flowControl(FlowControl flowControl) { initClientReference(flowControl.getClientId()); if (listener.compareAndSet(null, new GrpcQueryDispatcherListener(queryDispatcher, clientRef.get().toString(), wrappedQueryProviderInboundObserver, processingThreads))) { dispatcherListeners.put(clientRef.get(), listener.get()); } listener.get().addPermits(flowControl.getPermits()); } private void initClientReference(String clientId) { String clientStreamId = clientId + "." + UUID.randomUUID().toString(); if (clientRef.compareAndSet(null, new ClientStreamIdentification(context, clientStreamId))) { clientIdRegistry.register(clientStreamId, clientId, ClientIdRegistry.ConnectionType.QUERY); } clientIdRef.compareAndSet(null, clientId); } private void checkInitClient(String clientId, String componentName) { initClientReference(clientId); queryHandler.compareAndSet(null, new DirectQueryHandler(wrappedQueryProviderInboundObserver, clientRef.get(), componentName, clientId)); } @Override protected String sender() { return clientRef.toString(); } @Override public void onError(Throwable cause) { if (!ExceptionUtils.isCancelled(cause)) { logger.warn("{}: Error on connection from subscriber - {}", clientRef, cause.getMessage()); } cleanup(); } private void cleanup() { if (clientRef.get() != null) { String clientStreamId = clientRef.get().getClientStreamId(); String clientId = this.clientIdRef.get(); clientIdRegistry.unregister(clientStreamId, ClientIdRegistry.ConnectionType.QUERY); eventPublisher.publishEvent(new QueryHandlerDisconnected(context, clientId, clientStreamId)); } if (listener.get() != null) { dispatcherListeners.remove(clientRef.get()); listener.get().cancel(); } StreamObserverUtils.complete(inboundStreamObserver); } @Override public void onCompleted() { cleanup(); } }; } QueryService(Topology topology, QueryDispatcher queryDispatcher, ContextProvider contextProvider,
                        ClientIdRegistry clientIdRegistry,
                        ApplicationEventPublisher eventPublisher,
                        @Qualifier("queryInstructionAckSource")
                                InstructionAckSource<QueryProviderInbound> instructionAckSource); }
QueryService extends QueryServiceGrpc.QueryServiceImplBase implements AxonServerClientService { @Override public StreamObserver<QueryProviderOutbound> openStream( StreamObserver<QueryProviderInbound> inboundStreamObserver) { String context = contextProvider.getContext(); SendingStreamObserver<QueryProviderInbound> wrappedQueryProviderInboundObserver = new SendingStreamObserver<>( inboundStreamObserver); return new ReceivingStreamObserver<QueryProviderOutbound>(logger) { private final AtomicReference<String> clientIdRef = new AtomicReference<>(); private final AtomicReference<GrpcQueryDispatcherListener> listener = new AtomicReference<>(); private final AtomicReference<ClientStreamIdentification> clientRef = new AtomicReference<>(); private final AtomicReference<QueryHandler<QueryProviderInbound>> queryHandler = new AtomicReference<>(); @Override protected void consume(QueryProviderOutbound queryProviderOutbound) { switch (queryProviderOutbound.getRequestCase()) { case SUBSCRIBE: instructionAckSource.sendSuccessfulAck(queryProviderOutbound.getInstructionId(), wrappedQueryProviderInboundObserver); QuerySubscription subscription = queryProviderOutbound.getSubscribe(); checkInitClient(subscription.getClientId(), subscription.getComponentName()); String clientStreamId = clientRef.get().getClientStreamId(); logger.debug("{}: Subscribe Query {} for {}", context, subscription.getQuery(), subscription.getClientId()); SubscribeQuery subscribeQuery = new SubscribeQuery(context, clientStreamId, subscription, queryHandler.get()); eventPublisher.publishEvent( subscribeQuery); break; case UNSUBSCRIBE: instructionAckSource.sendSuccessfulAck(queryProviderOutbound.getInstructionId(), wrappedQueryProviderInboundObserver); if (clientRef.get() != null) { QuerySubscription unsubscribe = queryProviderOutbound.getUnsubscribe(); UnsubscribeQuery unsubscribeQuery = new UnsubscribeQuery(context, clientRef.get() .getClientStreamId(), unsubscribe, false); eventPublisher.publishEvent(unsubscribeQuery); } break; case FLOW_CONTROL: instructionAckSource.sendSuccessfulAck(queryProviderOutbound.getInstructionId(), wrappedQueryProviderInboundObserver); flowControl(queryProviderOutbound.getFlowControl()); break; case QUERY_RESPONSE: instructionAckSource.sendSuccessfulAck(queryProviderOutbound.getInstructionId(), wrappedQueryProviderInboundObserver); queryDispatcher.handleResponse(queryProviderOutbound.getQueryResponse(), clientRef.get().getClientStreamId(), clientIdRef.get(), false); break; case QUERY_COMPLETE: instructionAckSource.sendSuccessfulAck(queryProviderOutbound.getInstructionId(), wrappedQueryProviderInboundObserver); queryDispatcher.handleComplete(queryProviderOutbound.getQueryComplete().getRequestId(), clientRef.get().getClientStreamId(), clientIdRef.get(), false); break; case SUBSCRIPTION_QUERY_RESPONSE: instructionAckSource.sendSuccessfulAck(queryProviderOutbound.getInstructionId(), wrappedQueryProviderInboundObserver); SubscriptionQueryResponse response = queryProviderOutbound.getSubscriptionQueryResponse(); eventPublisher.publishEvent(new SubscriptionQueryResponseReceived(response, () -> wrappedQueryProviderInboundObserver .onNext(unsubscribeMessage(response.getSubscriptionIdentifier())))); break; case ACK: InstructionAck ack = queryProviderOutbound.getAck(); if (isUnsupportedInstructionErrorResult(ack)) { logger.warn("Unsupported instruction sent to the client {} of context {}.", clientRef.get().getClientStreamId(), context); } else { logger.trace("Received instruction ack from the client {} of context {}. Result {}.", clientRef.get().getClientStreamId(), context, ack); } break; default: instructionAckSource.sendUnsupportedInstruction(queryProviderOutbound.getInstructionId(), topology.getMe().getName(), wrappedQueryProviderInboundObserver); break; } } private void flowControl(FlowControl flowControl) { initClientReference(flowControl.getClientId()); if (listener.compareAndSet(null, new GrpcQueryDispatcherListener(queryDispatcher, clientRef.get().toString(), wrappedQueryProviderInboundObserver, processingThreads))) { dispatcherListeners.put(clientRef.get(), listener.get()); } listener.get().addPermits(flowControl.getPermits()); } private void initClientReference(String clientId) { String clientStreamId = clientId + "." + UUID.randomUUID().toString(); if (clientRef.compareAndSet(null, new ClientStreamIdentification(context, clientStreamId))) { clientIdRegistry.register(clientStreamId, clientId, ClientIdRegistry.ConnectionType.QUERY); } clientIdRef.compareAndSet(null, clientId); } private void checkInitClient(String clientId, String componentName) { initClientReference(clientId); queryHandler.compareAndSet(null, new DirectQueryHandler(wrappedQueryProviderInboundObserver, clientRef.get(), componentName, clientId)); } @Override protected String sender() { return clientRef.toString(); } @Override public void onError(Throwable cause) { if (!ExceptionUtils.isCancelled(cause)) { logger.warn("{}: Error on connection from subscriber - {}", clientRef, cause.getMessage()); } cleanup(); } private void cleanup() { if (clientRef.get() != null) { String clientStreamId = clientRef.get().getClientStreamId(); String clientId = this.clientIdRef.get(); clientIdRegistry.unregister(clientStreamId, ClientIdRegistry.ConnectionType.QUERY); eventPublisher.publishEvent(new QueryHandlerDisconnected(context, clientId, clientStreamId)); } if (listener.get() != null) { dispatcherListeners.remove(clientRef.get()); listener.get().cancel(); } StreamObserverUtils.complete(inboundStreamObserver); } @Override public void onCompleted() { cleanup(); } }; } QueryService(Topology topology, QueryDispatcher queryDispatcher, ContextProvider contextProvider,
                        ClientIdRegistry clientIdRegistry,
                        ApplicationEventPublisher eventPublisher,
                        @Qualifier("queryInstructionAckSource")
                                InstructionAckSource<QueryProviderInbound> instructionAckSource); @PreDestroy void cleanup(); @Override StreamObserver<QueryProviderOutbound> openStream(
            StreamObserver<QueryProviderInbound> inboundStreamObserver); @Override void query(QueryRequest request, StreamObserver<QueryResponse> responseObserver); @Override StreamObserver<SubscriptionQueryRequest> subscription(
            StreamObserver<SubscriptionQueryResponse> responseObserver); Set<GrpcQueryDispatcherListener> listeners(); void completeStream(String clientId, ClientStreamIdentification clientStreamIdentification); }
QueryService extends QueryServiceGrpc.QueryServiceImplBase implements AxonServerClientService { @Override public StreamObserver<QueryProviderOutbound> openStream( StreamObserver<QueryProviderInbound> inboundStreamObserver) { String context = contextProvider.getContext(); SendingStreamObserver<QueryProviderInbound> wrappedQueryProviderInboundObserver = new SendingStreamObserver<>( inboundStreamObserver); return new ReceivingStreamObserver<QueryProviderOutbound>(logger) { private final AtomicReference<String> clientIdRef = new AtomicReference<>(); private final AtomicReference<GrpcQueryDispatcherListener> listener = new AtomicReference<>(); private final AtomicReference<ClientStreamIdentification> clientRef = new AtomicReference<>(); private final AtomicReference<QueryHandler<QueryProviderInbound>> queryHandler = new AtomicReference<>(); @Override protected void consume(QueryProviderOutbound queryProviderOutbound) { switch (queryProviderOutbound.getRequestCase()) { case SUBSCRIBE: instructionAckSource.sendSuccessfulAck(queryProviderOutbound.getInstructionId(), wrappedQueryProviderInboundObserver); QuerySubscription subscription = queryProviderOutbound.getSubscribe(); checkInitClient(subscription.getClientId(), subscription.getComponentName()); String clientStreamId = clientRef.get().getClientStreamId(); logger.debug("{}: Subscribe Query {} for {}", context, subscription.getQuery(), subscription.getClientId()); SubscribeQuery subscribeQuery = new SubscribeQuery(context, clientStreamId, subscription, queryHandler.get()); eventPublisher.publishEvent( subscribeQuery); break; case UNSUBSCRIBE: instructionAckSource.sendSuccessfulAck(queryProviderOutbound.getInstructionId(), wrappedQueryProviderInboundObserver); if (clientRef.get() != null) { QuerySubscription unsubscribe = queryProviderOutbound.getUnsubscribe(); UnsubscribeQuery unsubscribeQuery = new UnsubscribeQuery(context, clientRef.get() .getClientStreamId(), unsubscribe, false); eventPublisher.publishEvent(unsubscribeQuery); } break; case FLOW_CONTROL: instructionAckSource.sendSuccessfulAck(queryProviderOutbound.getInstructionId(), wrappedQueryProviderInboundObserver); flowControl(queryProviderOutbound.getFlowControl()); break; case QUERY_RESPONSE: instructionAckSource.sendSuccessfulAck(queryProviderOutbound.getInstructionId(), wrappedQueryProviderInboundObserver); queryDispatcher.handleResponse(queryProviderOutbound.getQueryResponse(), clientRef.get().getClientStreamId(), clientIdRef.get(), false); break; case QUERY_COMPLETE: instructionAckSource.sendSuccessfulAck(queryProviderOutbound.getInstructionId(), wrappedQueryProviderInboundObserver); queryDispatcher.handleComplete(queryProviderOutbound.getQueryComplete().getRequestId(), clientRef.get().getClientStreamId(), clientIdRef.get(), false); break; case SUBSCRIPTION_QUERY_RESPONSE: instructionAckSource.sendSuccessfulAck(queryProviderOutbound.getInstructionId(), wrappedQueryProviderInboundObserver); SubscriptionQueryResponse response = queryProviderOutbound.getSubscriptionQueryResponse(); eventPublisher.publishEvent(new SubscriptionQueryResponseReceived(response, () -> wrappedQueryProviderInboundObserver .onNext(unsubscribeMessage(response.getSubscriptionIdentifier())))); break; case ACK: InstructionAck ack = queryProviderOutbound.getAck(); if (isUnsupportedInstructionErrorResult(ack)) { logger.warn("Unsupported instruction sent to the client {} of context {}.", clientRef.get().getClientStreamId(), context); } else { logger.trace("Received instruction ack from the client {} of context {}. Result {}.", clientRef.get().getClientStreamId(), context, ack); } break; default: instructionAckSource.sendUnsupportedInstruction(queryProviderOutbound.getInstructionId(), topology.getMe().getName(), wrappedQueryProviderInboundObserver); break; } } private void flowControl(FlowControl flowControl) { initClientReference(flowControl.getClientId()); if (listener.compareAndSet(null, new GrpcQueryDispatcherListener(queryDispatcher, clientRef.get().toString(), wrappedQueryProviderInboundObserver, processingThreads))) { dispatcherListeners.put(clientRef.get(), listener.get()); } listener.get().addPermits(flowControl.getPermits()); } private void initClientReference(String clientId) { String clientStreamId = clientId + "." + UUID.randomUUID().toString(); if (clientRef.compareAndSet(null, new ClientStreamIdentification(context, clientStreamId))) { clientIdRegistry.register(clientStreamId, clientId, ClientIdRegistry.ConnectionType.QUERY); } clientIdRef.compareAndSet(null, clientId); } private void checkInitClient(String clientId, String componentName) { initClientReference(clientId); queryHandler.compareAndSet(null, new DirectQueryHandler(wrappedQueryProviderInboundObserver, clientRef.get(), componentName, clientId)); } @Override protected String sender() { return clientRef.toString(); } @Override public void onError(Throwable cause) { if (!ExceptionUtils.isCancelled(cause)) { logger.warn("{}: Error on connection from subscriber - {}", clientRef, cause.getMessage()); } cleanup(); } private void cleanup() { if (clientRef.get() != null) { String clientStreamId = clientRef.get().getClientStreamId(); String clientId = this.clientIdRef.get(); clientIdRegistry.unregister(clientStreamId, ClientIdRegistry.ConnectionType.QUERY); eventPublisher.publishEvent(new QueryHandlerDisconnected(context, clientId, clientStreamId)); } if (listener.get() != null) { dispatcherListeners.remove(clientRef.get()); listener.get().cancel(); } StreamObserverUtils.complete(inboundStreamObserver); } @Override public void onCompleted() { cleanup(); } }; } QueryService(Topology topology, QueryDispatcher queryDispatcher, ContextProvider contextProvider,
                        ClientIdRegistry clientIdRegistry,
                        ApplicationEventPublisher eventPublisher,
                        @Qualifier("queryInstructionAckSource")
                                InstructionAckSource<QueryProviderInbound> instructionAckSource); @PreDestroy void cleanup(); @Override StreamObserver<QueryProviderOutbound> openStream(
            StreamObserver<QueryProviderInbound> inboundStreamObserver); @Override void query(QueryRequest request, StreamObserver<QueryResponse> responseObserver); @Override StreamObserver<SubscriptionQueryRequest> subscription(
            StreamObserver<SubscriptionQueryResponse> responseObserver); Set<GrpcQueryDispatcherListener> listeners(); void completeStream(String clientId, ClientStreamIdentification clientStreamIdentification);  }