Skip to content

Commit a396f24

Browse files
dpcollins-googledpcollins-google
authored
feat: Add the ability to use automatic subscriber assignment to the subscriber settings (#163)
* [feat] Add the ability to use automatic subscriber assignment to the subscriber settings. Also disable clirr as it fires on this PR. We are not providing interface stability guarantees. * No-op commit to update github. * Undo sample changes. * Revert lastt samples change.
1 parent 8effb81 commit a396f24

File tree

3 files changed

+161
-13
lines changed

3 files changed

+161
-13
lines changed

‎google-cloud-pubsublite/pom.xml‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -215,6 +215,21 @@
215215
<groupId>org.codehaus.mojo</groupId>
216216
<artifactId>flatten-maven-plugin</artifactId>
217217
</plugin>
218+
<plugin>
219+
<!--TODO: Remove before GA. -->
220+
<groupId>org.codehaus.mojo</groupId>
221+
<artifactId>clirr-maven-plugin</artifactId>
222+
<configuration>
223+
<failOnError>false</failOnError>
224+
</configuration>
225+
<executions>
226+
<execution>
227+
<goals>
228+
<goal>check</goal>
229+
</goals>
230+
</execution>
231+
</executions>
232+
</plugin>
218233
</plugins>
219234
</build>
220235
</project>

‎google-cloud-pubsublite/src/main/java/com/google/cloud/pubsublite/cloudpubsub/SubscriberSettings.java‎

Lines changed: 43 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -24,16 +24,20 @@
2424
import com.google.cloud.pubsublite.SubscriptionPath;
2525
import com.google.cloud.pubsublite.SubscriptionPaths;
2626
import com.google.cloud.pubsublite.cloudpubsub.internal.AckSetTrackerImpl;
27+
import com.google.cloud.pubsublite.cloudpubsub.internal.AssigningSubscriber;
2728
import com.google.cloud.pubsublite.cloudpubsub.internal.MultiPartitionSubscriber;
29+
import com.google.cloud.pubsublite.cloudpubsub.internal.PartitionSubscriberFactory;
2830
import com.google.cloud.pubsublite.cloudpubsub.internal.SinglePartitionSubscriber;
2931
import com.google.cloud.pubsublite.internal.Preconditions;
32+
import com.google.cloud.pubsublite.internal.wire.AssignerBuilder;
33+
import com.google.cloud.pubsublite.internal.wire.AssignerFactory;
3034
import com.google.cloud.pubsublite.internal.wire.CommitterBuilder;
3135
import com.google.cloud.pubsublite.internal.wire.PubsubContext;
3236
import com.google.cloud.pubsublite.internal.wire.PubsubContext.Framework;
3337
import com.google.cloud.pubsublite.internal.wire.SubscriberBuilder;
3438
import com.google.cloud.pubsublite.proto.CursorServiceGrpc;
39+
import com.google.cloud.pubsublite.proto.PartitionAssignmentServiceGrpc.PartitionAssignmentServiceStub;
3540
import com.google.cloud.pubsublite.proto.SubscriberServiceGrpc;
36-
import com.google.common.collect.ImmutableList;
3741
import com.google.pubsub.v1.PubsubMessage;
3842
import io.grpc.StatusException;
3943
import java.util.ArrayList;
@@ -51,17 +55,21 @@ public abstract class SubscriberSettings {
5155

5256
abstract SubscriptionPath subscriptionPath();
5357

54-
abstract ImmutableList<Partition> partitions();
55-
5658
abstract FlowControlSettings perPartitionFlowControlSettings();
5759

5860
// Optional parameters.
61+
62+
// If set, disables auto-assignment.
63+
abstract Optional<List<Partition>> partitions();
64+
5965
abstract Optional<MessageTransformer<SequencedMessage, PubsubMessage>> transformer();
6066

6167
abstract Optional<SubscriberServiceGrpc.SubscriberServiceStub> subscriberServiceStub();
6268

6369
abstract Optional<CursorServiceGrpc.CursorServiceStub> cursorServiceStub();
6470

71+
abstract Optional<PartitionAssignmentServiceStub> assignmentServiceStub();
72+
6573
abstract Optional<NackHandler> nackHandler();
6674

6775
public static Builder newBuilder() {
@@ -76,11 +84,12 @@ public abstract static class Builder {
7684

7785
public abstract Builder setSubscriptionPath(SubscriptionPath path);
7886

79-
public abstract Builder setPartitions(List<Partition> partition);
80-
8187
public abstract Builder setPerPartitionFlowControlSettings(FlowControlSettings settings);
8288

8389
// Optional parameters.
90+
/** If set, disables auto-assignment. */
91+
public abstract Builder setPartitions(List<Partition> partition);
92+
8493
public abstract Builder setTransformer(
8594
MessageTransformer<SequencedMessage, PubsubMessage> transformer);
8695

@@ -89,14 +98,17 @@ public abstract Builder setSubscriberServiceStub(
8998

9099
public abstract Builder setCursorServiceStub(CursorServiceGrpc.CursorServiceStub stub);
91100

101+
public abstract Builder setAssignmentServiceStub(PartitionAssignmentServiceStub stub);
102+
92103
public abstract Builder setNackHandler(NackHandler nackHandler);
93104

94105
abstract SubscriberSettings autoBuild();
95106

96107
public SubscriberSettings build() throws StatusException {
97108
SubscriberSettings settings = autoBuild();
98109
Preconditions.checkArgument(
99-
!settings.partitions().isEmpty(), "Must provide at least one partition.");
110+
!settings.partitions().isPresent() || !settings.partitions().get().isEmpty(),
111+
"Must provide at least one partition if setting partitions explicitly.");
100112
SubscriptionPaths.check(settings.subscriptionPath());
101113
return settings;
102114
}
@@ -113,18 +125,36 @@ Subscriber instantiate() throws StatusException {
113125
wireCommitterBuilder.setSubscriptionPath(subscriptionPath());
114126
cursorServiceStub().ifPresent(wireCommitterBuilder::setCursorStub);
115127

116-
List<Subscriber> perPartitionSubscribers = new ArrayList<>();
117-
for (Partition partition : partitions()) {
118-
wireSubscriberBuilder.setPartition(partition);
119-
wireCommitterBuilder.setPartition(partition);
120-
perPartitionSubscribers.add(
121-
new SinglePartitionSubscriber(
128+
PartitionSubscriberFactory partitionSubscriberFactory =
129+
partition -> {
130+
wireSubscriberBuilder.setPartition(partition);
131+
wireCommitterBuilder.setPartition(partition);
132+
return new SinglePartitionSubscriber(
122133
receiver(),
123134
transformer().orElse(MessageTransforms.toCpsSubscribeTransformer()),
124135
new AckSetTrackerImpl(wireCommitterBuilder.build()),
125136
nackHandler().orElse(new NackHandler() {}),
126137
messageConsumer -> wireSubscriberBuilder.setMessageConsumer(messageConsumer).build(),
127-
perPartitionFlowControlSettings()));
138+
perPartitionFlowControlSettings());
139+
};
140+
141+
if (!partitions().isPresent()) {
142+
AssignerBuilder.Builder assignerBuilder = AssignerBuilder.newBuilder();
143+
assignerBuilder.setSubscriptionPath(subscriptionPath());
144+
assignmentServiceStub().ifPresent(assignerBuilder::setAssignmentStub);
145+
AssignerFactory assignerFactory =
146+
receiver -> {
147+
assignerBuilder.setReceiver(receiver);
148+
return assignerBuilder.build();
149+
};
150+
return new AssigningSubscriber(partitionSubscriberFactory, assignerFactory);
151+
}
152+
153+
List<Subscriber> perPartitionSubscribers = new ArrayList<>();
154+
for (Partition partition : partitions().get()) {
155+
wireSubscriberBuilder.setPartition(partition);
156+
wireCommitterBuilder.setPartition(partition);
157+
perPartitionSubscribers.add(partitionSubscriberFactory.New(partition));
128158
}
129159
return MultiPartitionSubscriber.of(perPartitionSubscribers);
130160
}

‎google-cloud-pubsublite/src/main/java/com/google/cloud/pubsublite/internal/wire/AssignerBuilder.java‎

Lines changed: 103 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,103 @@
1+
/*
2+
* Copyright 2020 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.google.cloud.pubsublite.internal.wire;
18+
19+
import com.google.auto.value.AutoValue;
20+
import com.google.cloud.pubsublite.Endpoints;
21+
import com.google.cloud.pubsublite.Stubs;
22+
import com.google.cloud.pubsublite.SubscriptionPath;
23+
import com.google.cloud.pubsublite.SubscriptionPaths;
24+
import com.google.cloud.pubsublite.proto.InitialPartitionAssignmentRequest;
25+
import com.google.cloud.pubsublite.proto.PartitionAssignmentServiceGrpc;
26+
import com.google.cloud.pubsublite.proto.PartitionAssignmentServiceGrpc.PartitionAssignmentServiceStub;
27+
import com.google.common.flogger.GoogleLogger;
28+
import com.google.protobuf.ByteString;
29+
import io.grpc.Status;
30+
import io.grpc.StatusException;
31+
import java.io.IOException;
32+
import java.nio.ByteBuffer;
33+
import java.util.Optional;
34+
import java.util.UUID;
35+
36+
@AutoValue
37+
public abstract class AssignerBuilder {
38+
private static final GoogleLogger logger = GoogleLogger.forEnclosingClass();
39+
// Required parameters.
40+
abstract SubscriptionPath subscriptionPath();
41+
42+
abstract PartitionAssignmentReceiver receiver();
43+
44+
// Optional parameters.
45+
abstract Optional<PartitionAssignmentServiceStub> assignmentStub();
46+
47+
public static Builder newBuilder() {
48+
return new AutoValue_AssignerBuilder.Builder();
49+
}
50+
51+
@AutoValue.Builder
52+
public abstract static class Builder {
53+
// Required parameters.
54+
public abstract Builder setSubscriptionPath(SubscriptionPath path);
55+
56+
public abstract Builder setReceiver(PartitionAssignmentReceiver receiver);
57+
58+
// Optional parameters.
59+
public abstract Builder setAssignmentStub(PartitionAssignmentServiceStub stub);
60+
61+
abstract AssignerBuilder autoBuild();
62+
63+
@SuppressWarnings("CheckReturnValue")
64+
public Assigner build() throws StatusException {
65+
AssignerBuilder builder = autoBuild();
66+
SubscriptionPaths.check(builder.subscriptionPath());
67+
68+
PartitionAssignmentServiceStub stub;
69+
if (builder.assignmentStub().isPresent()) {
70+
stub = builder.assignmentStub().get();
71+
} else {
72+
try {
73+
stub =
74+
Stubs.defaultStub(
75+
Endpoints.regionalEndpoint(
76+
SubscriptionPaths.getZone(builder.subscriptionPath()).region()),
77+
PartitionAssignmentServiceGrpc::newStub);
78+
} catch (IOException e) {
79+
throw Status.INTERNAL
80+
.withCause(e)
81+
.withDescription("Creating assigner stub failed.")
82+
.asException();
83+
}
84+
}
85+
86+
UUID uuid = UUID.randomUUID();
87+
ByteBuffer uuidBuffer = ByteBuffer.allocate(16);
88+
uuidBuffer.putLong(uuid.getMostSignificantBits());
89+
uuidBuffer.putLong(uuid.getLeastSignificantBits());
90+
logger.atInfo().log(
91+
"Subscription %s using UUID %s for assignment.",
92+
builder.subscriptionPath().value(), uuid);
93+
94+
InitialPartitionAssignmentRequest initial =
95+
InitialPartitionAssignmentRequest.newBuilder()
96+
.setSubscription(builder.subscriptionPath().value())
97+
.setClientId(ByteString.copyFrom(uuidBuffer.array()))
98+
.build();
99+
return new AssignerImpl(
100+
stub, new ConnectedAssignerImpl.Factory(), initial, builder.receiver());
101+
}
102+
}
103+
}

0 commit comments

Comments
 (0)