From 6dd5efe4ad4b88afcc22ff9184f2210735d71b90 Mon Sep 17 00:00:00 2001
From: Oscar Scholten <oscar.scholten@axoniq.io>
Date: Fri, 24 Nov 2017 17:29:27 +0100
Subject: [PATCH] Return Registration object when registering interceptor

---
 .../axonframework/queryhandling/SimpleQueryBus.java  | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/core/src/main/java/org/axonframework/queryhandling/SimpleQueryBus.java b/core/src/main/java/org/axonframework/queryhandling/SimpleQueryBus.java
index 73386927dc..44b429efe5 100644
--- a/core/src/main/java/org/axonframework/queryhandling/SimpleQueryBus.java
+++ b/core/src/main/java/org/axonframework/queryhandling/SimpleQueryBus.java
@@ -193,9 +193,11 @@ private boolean unsubscribe(QueryDefinition registrationKey, MessageHandler<? su
      * respective handlers. The interceptor is invoked separately for each handler instance (in a separate unit of work).
      *
      * @param interceptor the interceptor to invoke before passing a Query to the handler
+     * @return handle to unregister the interceptor
      */
-    public void registerHandlerInterceptor(MessageHandlerInterceptor<QueryMessage<?, ?>> interceptor) {
-        this.handlerInterceptors.add(interceptor);
+    public Registration registerHandlerInterceptor(MessageHandlerInterceptor<QueryMessage<?, ?>> interceptor) {
+        handlerInterceptors.add(interceptor);
+        return () -> handlerInterceptors.remove(interceptor);
     }
 
     /**
@@ -203,9 +205,11 @@ public void registerHandlerInterceptor(MessageHandlerInterceptor<QueryMessage<?,
      * once, regardless of the type of query (point-to-point or scatter-gather) executed.
      *
      * @param interceptor the interceptor to invoke when sending a Query
+     * @return handle to unregister the interceptor
      */
-    public void registerDispatchInterceptor(MessageDispatchInterceptor<QueryMessage<?, ?>> interceptor) {
-        this.dispatchInterceptors.add(interceptor);
+    public Registration registerDispatchInterceptor(MessageDispatchInterceptor<QueryMessage<?, ?>> interceptor) {
+        dispatchInterceptors.add(interceptor);
+        return () -> dispatchInterceptors.remove(interceptor);
     }
 
     private <Q, R> Set<MessageHandler<? super QueryMessage<?, ?>>> getHandlersForMessage(QueryMessage<Q, R> queryMessage) {