cadonna commented on code in PR #17795:
URL: https://github.com/apache/kafka/pull/17795#discussion_r1858311885


##########
clients/src/main/java/org/apache/kafka/clients/consumer/internals/StreamsAssignmentInterface.java:
##########
@@ -281,17 +300,173 @@ public String toString() {
         }
     }
 
+    private final BlockingQueue<BackgroundEvent> onCallbackRequests = new 
LinkedBlockingQueue<>();
+
+    private ApplicationEventHandler applicationEventHandler = null;
+
+    private Optional<Function<Set<StreamsAssignmentInterface.TaskId>, 
Optional<Exception>>> onTasksRevokedCallback = null;
+    private Optional<Function<Assignment, Optional<Exception>>> 
onTasksAssignedCallback = null;
+    private Optional<Supplier<Optional<Exception>>> onAllTasksLostCallback = 
null;
+
+    private final StreamsRebalanceEventProcessor 
streamsRebalanceEventProcessor;
+
+    private class StreamsRebalanceEventProcessor implements 
EventProcessor<BackgroundEvent> {
+
+        @Override
+        public void process(final BackgroundEvent event) {
+            switch (event.type()) {
+                case ERROR:
+                    process((ErrorEvent) event);
+                    break;
+
+                case STREAMS_ON_TASKS_REVOKED_CALLBACK_NEEDED:
+                    process((StreamsOnTasksRevokedCallbackNeededEvent) event);
+                    break;
+
+                case STREAMS_ON_TASKS_ASSIGNED_CALLBACK_NEEDED:
+                    process((StreamsOnTasksAssignedCallbackNeededEvent) event);
+                    break;
+
+                case STREAMS_ON_ALL_TASKS_LOST_CALLBACK_NEEDED:
+                    process((StreamsOnAllTasksLostCallbackNeededEvent) event);
+                    break;
+
+                default:
+                    throw new IllegalArgumentException("Background event type 
" + event.type() + " was not expected");
+
+            }
+        }
+
+        private void process(final ErrorEvent event) {
+            throw event.error();
+        }
+
+        private void process(final StreamsOnTasksRevokedCallbackNeededEvent 
event) {
+            StreamsOnTasksRevokedCallbackCompletedEvent invokedEvent = 
invokeOnTasksRevokedCallback(event.activeTasksToRevoke(), event.future());
+            applicationEventHandler.add(invokedEvent);
+            if (invokedEvent.error().isPresent()) {
+                throw invokedEvent.error().get();
+            }
+        }
+
+        private void process(final StreamsOnTasksAssignedCallbackNeededEvent 
event) {
+            StreamsOnTasksAssignedCallbackCompletedEvent invokedEvent = 
invokeOnTasksAssignedCallback(event.assignment(), event.future());
+            applicationEventHandler.add(invokedEvent);
+            if (invokedEvent.error().isPresent()) {
+                throw invokedEvent.error().get();
+            }
+        }
+
+        private void process(final StreamsOnAllTasksLostCallbackNeededEvent 
event) {
+            StreamsOnAllTasksLostCallbackCompletedEvent invokedEvent = 
invokeOnAllTasksLostCallback(event.future());
+            applicationEventHandler.add(invokedEvent);
+            if (invokedEvent.error().isPresent()) {
+                throw invokedEvent.error().get();
+            }
+        }
+
+        private StreamsOnTasksRevokedCallbackCompletedEvent 
invokeOnTasksRevokedCallback(final Set<StreamsAssignmentInterface.TaskId> 
activeTasksToRevoke,
+                                                                               
           final CompletableFuture<Void> future) {
+            final Optional<Exception> exceptionFromCallback = 
onTasksRevokedCallback
+                .orElseThrow(() -> new IllegalStateException("No tasks 
assignment callback set!")).apply(activeTasksToRevoke);
+
+            return exceptionFromCallback
+                .map(exception ->
+                    new StreamsOnTasksRevokedCallbackCompletedEvent(
+                        future,
+                        
Optional.of(ConsumerUtils.maybeWrapAsKafkaException(exception, "Task revocation 
callback throws an error"))
+                    ))
+                .orElseGet(() -> new 
StreamsOnTasksRevokedCallbackCompletedEvent(future, Optional.empty()));
+        }
+
+        private StreamsOnTasksAssignedCallbackCompletedEvent 
invokeOnTasksAssignedCallback(final StreamsAssignmentInterface.Assignment 
assignment,
+                                                                               
            final CompletableFuture<Void> future) {
+            Optional<KafkaException> error = Optional.empty();
+            // ToDo: Can we avoid the following check?
+            if (!assignment.equals(reconciledAssignment.get())) {
+
+                final Optional<Exception> exceptionFromCallback = 
onTasksAssignedCallback
+                    .orElseThrow(() -> new IllegalStateException("No tasks 
assignment callback set!")).apply(assignment);
+
+                if (exceptionFromCallback.isPresent()) {
+                    error = 
Optional.of(ConsumerUtils.maybeWrapAsKafkaException(exceptionFromCallback.get(),
 "Task assignment callback throws an error"));
+                } else {
+                    reconciledAssignment.set(assignment);
+                }
+            }
+            return new StreamsOnTasksAssignedCallbackCompletedEvent(future, 
error);
+        }
+
+        private StreamsOnAllTasksLostCallbackCompletedEvent 
invokeOnAllTasksLostCallback(final CompletableFuture<Void> future) {
+            final Optional<Exception> exceptionFromCallback = 
onAllTasksLostCallback
+                .orElseThrow(() -> new IllegalStateException("No tasks 
assignment callback set!")).get();
+
+            final Optional<KafkaException> error;
+
+            if (exceptionFromCallback.isPresent()) {
+                error = 
Optional.of(ConsumerUtils.maybeWrapAsKafkaException(exceptionFromCallback.get(),
 "Task assignment callback throws an error"));
+            } else {
+                error = Optional.empty();
+                reconciledAssignment.set(Assignment.EMPTY);
+            }
+
+            return new StreamsOnAllTasksLostCallbackCompletedEvent(future, 
error);
+        }
+    }
+
     public StreamsAssignmentInterface(UUID processId,
                                       Optional<HostInfo> endpoint,
                                       Map<String, Subtopology> subtopologyMap,
-                                      Map<String, String> clientTags
-    ) {
+                                      Map<String, String> clientTags) {
         this.processId = processId;
         this.endpoint = endpoint;
         this.subtopologyMap = subtopologyMap;
         this.taskLags = new HashMap<>();
         this.shutdownRequested = new AtomicBoolean(false);
         this.clientTags = clientTags;
+        this.streamsRebalanceEventProcessor = new 
StreamsRebalanceEventProcessor();
+    }
+
+    public void setOnTasksRevokedCallback(final 
Function<Set<StreamsAssignmentInterface.TaskId>, Optional<Exception>> 
onTasksRevokedCallback) {
+        this.onTasksRevokedCallback = 
Optional.ofNullable(onTasksRevokedCallback);
+    }
+
+    public void setOnTasksAssignedCallback(final Function<Assignment, 
Optional<Exception>> onTasksAssignedCallback) {
+        this.onTasksAssignedCallback = 
Optional.ofNullable(onTasksAssignedCallback);
+    }
+
+    public void setOnAllTasksLostCallback(final Supplier<Optional<Exception>> 
onAllTasksLostCallback) {
+        this.onAllTasksLostCallback = 
Optional.ofNullable(onAllTasksLostCallback);
+    }
+
+    public void setApplicationEventHandler(final ApplicationEventHandler 
applicationEventHandler) {
+        this.applicationEventHandler = applicationEventHandler;
+    }
+
+    public CompletableFuture<Void> 
requestOnTasksAssignedCallbackInvocation(final Assignment assignment) {
+        final StreamsOnTasksAssignedCallbackNeededEvent 
onTasksAssignedCallbackNeededEvent = new 
StreamsOnTasksAssignedCallbackNeededEvent(assignment);
+        onCallbackRequests.add(onTasksAssignedCallbackNeededEvent);
+        return onTasksAssignedCallbackNeededEvent.future();
+    }
+
+    public CompletableFuture<Void> 
requestOnTasksRevokedCallbackInvocation(final 
Set<StreamsAssignmentInterface.TaskId> activeTasksToRevoke) {
+        final StreamsOnTasksRevokedCallbackNeededEvent 
onTasksRevokedCallbackNeededEvent = new 
StreamsOnTasksRevokedCallbackNeededEvent(activeTasksToRevoke);
+        onCallbackRequests.add(onTasksRevokedCallbackNeededEvent);
+        return onTasksRevokedCallbackNeededEvent.future();
+    }
+
+    public CompletableFuture<Void> requestOnAllTasksLostCallbackInvocation() {

Review Comment:
   Basically, we use the futures only to bind code to them so that this code is 
executed when they complete. We pass the future from the background thread to 
the application thread which does something. Then when this something is done, 
we pass the future back to the background thread where it is completed and the 
code is executed. It is really just a way to pass around code and execute it at 
the right moment.   



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: jira-unsubscr...@kafka.apache.org

For queries about this service, please contact Infrastructure at:
us...@infra.apache.org

Reply via email to