1
1
/*
2
- * Copyright 2019-2023 the original author or authors.
2
+ * Copyright 2019-2024 the original author or authors.
3
3
*
4
4
* Licensed under the Apache License, Version 2.0 (the "License");
5
5
* you may not use this file except in compliance with the License.
16
16
17
17
package org .springframework .kafka .listener ;
18
18
19
+ import java .util .ArrayList ;
19
20
import java .util .Collection ;
20
21
import java .util .Collections ;
21
22
import java .util .LinkedList ;
22
23
import java .util .List ;
23
24
import java .util .Map ;
24
25
import java .util .concurrent .ConcurrentHashMap ;
25
-
26
26
import org .apache .kafka .common .TopicPartition ;
27
-
28
27
import org .springframework .lang .Nullable ;
29
28
30
29
/**
@@ -41,6 +40,8 @@ public abstract class AbstractConsumerSeekAware implements ConsumerSeekAware {
41
40
private final Map <Thread , ConsumerSeekCallback > callbackForThread = new ConcurrentHashMap <>();
42
41
43
42
private final Map <TopicPartition , ConsumerSeekCallback > callbacks = new ConcurrentHashMap <>();
43
+ // [Suggestion]
44
+ private final Map <TopicPartition , List <ConsumerSeekCallback >> callbacksV2 = new ConcurrentHashMap <>();
44
45
45
46
private final Map <ConsumerSeekCallback , List <TopicPartition >> callbacksToTopic = new ConcurrentHashMap <>();
46
47
@@ -60,6 +61,17 @@ public void onPartitionsAssigned(Map<TopicPartition, Long> assignments, Consumer
60
61
}
61
62
}
62
63
64
+ // [Suggestion]
65
+ public void onPartitionsAssignedV2 (Map <TopicPartition , Long > assignments , ConsumerSeekCallback callback ) {
66
+ ConsumerSeekCallback threadCallback = this .callbackForThread .get (Thread .currentThread ());
67
+ if (threadCallback != null ) {
68
+ assignments .keySet ().forEach (tp -> {
69
+ this .callbacksV2 .computeIfAbsent (tp , key -> new ArrayList <>()).add (threadCallback );
70
+ this .callbacksToTopic .computeIfAbsent (threadCallback , key -> new LinkedList <>()).add (tp );
71
+ });
72
+ }
73
+ }
74
+
63
75
@ Override
64
76
public void onPartitionsRevoked (Collection <TopicPartition > partitions ) {
65
77
partitions .forEach (tp -> {
@@ -76,6 +88,24 @@ public void onPartitionsRevoked(Collection<TopicPartition> partitions) {
76
88
});
77
89
}
78
90
91
+ // [Suggestion]
92
+ public void onPartitionsRevokedV2 (Collection <TopicPartition > partitions ) {
93
+ partitions .forEach (tp -> {
94
+ List <ConsumerSeekCallback > removed = this .callbacksV2 .remove (tp );
95
+ if (removed != null && !removed .isEmpty ()) {
96
+ removed .forEach (cb -> {
97
+ List <TopicPartition > topics = this .callbacksToTopic .get (cb );
98
+ if (topics != null ) {
99
+ topics .remove (tp );
100
+ if (topics .isEmpty ()) {
101
+ this .callbacksToTopic .remove (cb );
102
+ }
103
+ }
104
+ });
105
+ }
106
+ });
107
+ }
108
+
79
109
@ Override
80
110
public void unregisterSeekCallback () {
81
111
this .callbackForThread .remove (Thread .currentThread ());
@@ -91,6 +121,11 @@ protected ConsumerSeekCallback getSeekCallbackFor(TopicPartition topicPartition)
91
121
return this .callbacks .get (topicPartition );
92
122
}
93
123
124
+ // [Suggestion]
125
+ protected List <ConsumerSeekCallback > getSeekCallbackForV2 (TopicPartition topicPartition ) {
126
+ return this .callbacksV2 .get (topicPartition );
127
+ }
128
+
94
129
/**
95
130
* The map of callbacks for all currently assigned partitions.
96
131
* @return the map.
@@ -99,6 +134,11 @@ protected Map<TopicPartition, ConsumerSeekCallback> getSeekCallbacks() {
99
134
return Collections .unmodifiableMap (this .callbacks );
100
135
}
101
136
137
+ // [Suggestion]
138
+ protected Map <TopicPartition , List <ConsumerSeekCallback >> getSeekCallbacksV2 () {
139
+ return Collections .unmodifiableMap (this .callbacksV2 );
140
+ }
141
+
102
142
/**
103
143
* Return the currently registered callbacks and their associated {@link TopicPartition}(s).
104
144
* @return the map of callbacks and partitions.
0 commit comments