Add a context that allows to register RequestSateProviders
It's possible to register multiple RequestSateProvider's, so that
setting a RequestSateProvider doesn't override any RequestSateProvider
that has been set before.
The registered RequestStateProviders are stored in ThreadLocal so that
they can be accessed during the request execution.
On close the RequestStateProviders that have been registered by the
RequestStateContext instance are removed from ThreadLocal.
Signed-off-by: Edwin Kempin <ekempin@google.com>
Change-Id: Ie9dc1d14ed88bc3092d5848e9557dfa82ba0d533
diff --git a/java/com/google/gerrit/server/cancellation/RequestStateContext.java b/java/com/google/gerrit/server/cancellation/RequestStateContext.java
new file mode 100644
index 0000000..1fa2e03
--- /dev/null
+++ b/java/com/google/gerrit/server/cancellation/RequestStateContext.java
@@ -0,0 +1,117 @@
+// Copyright (C) 2021 The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.google.gerrit.server.cancellation;
+
+import com.google.common.collect.ImmutableSet;
+import java.util.HashSet;
+import java.util.Set;
+
+/**
+ * Context that allows to register {@link RequestStateProvider}s.
+ *
+ * <p>The registered {@link RequestStateProvider}s are stored in {@link ThreadLocal} so that they
+ * can be accessed during the request execution (via {@link #getRequestStateProviders()}.
+ *
+ * <p>On {@link #close()} the {@link RequestStateProvider}s that have been registered by this {@code
+ * RequestStateContext} instance are removed from {@link ThreadLocal}.
+ *
+ * <p>Nesting {@code RequestStateContext}s is possible.
+ *
+ * <p>Currently there is no logic to automatically copy the {@link RequestStateContext} to
+ * background threads, but implementing this may be considered in the future. This means that by
+ * default we only support cancellation of the main thread, but not of background threads. That's
+ * fine as all significant work is being done in the main thread.
+ *
+ * <p>{@link com.google.gerrit.server.util.RequestContext} is also a context that is available for
+ * the time of the request, but it is not suitable to manage registrations of {@link
+ * RequestStateProvider}s. Hence {@link RequestStateProvider} registrations are managed by a
+ * separate context, which is this class, {@link RequestStateContext}:
+ *
+ * <ul>
+ * <li>{@link com.google.gerrit.server.util.RequestContext} is an interface that has many
+ * implementations and hence cannot manage a {@link ThreadLocal} state.
+ * <li>{@link com.google.gerrit.server.util.RequestContext} is not an {@link AutoCloseable} and
+ * hence cannot cleanup any {@link ThreadLocal} state on close (turning it into an {@link
+ * AutoCloseable} would require a large refactoring).
+ * <li>Despite the name {@link com.google.gerrit.server.util.RequestContext} is not only used for
+ * requests scopes but also for other scopes that are not a request (e.g. plugin invocations,
+ * email sending, manual scopes).
+ * <li>{@link com.google.gerrit.server.util.RequestContext} is not copied to background and should
+ * not be, but for {@link RequestStateContext} we may consider doing this in the future.
+ * </ul>
+ */
+public class RequestStateContext implements AutoCloseable {
+ /** The {@link RequestStateProvider}s that have been registered for the thread. */
+ private static final ThreadLocal<Set<RequestStateProvider>> threadLocalRequestStateProviders =
+ new ThreadLocal<>();
+
+ /** Returns the {@link RequestStateProvider}s that have been registered for the thread. */
+ static ImmutableSet<RequestStateProvider> getRequestStateProviders() {
+ if (threadLocalRequestStateProviders.get() == null) {
+ return ImmutableSet.of();
+ }
+ return ImmutableSet.copyOf(threadLocalRequestStateProviders.get());
+ }
+
+ /** Opens a {@code RequestStateContext}. */
+ public static RequestStateContext open() {
+ return new RequestStateContext();
+ }
+
+ /**
+ * The {@link RequestStateProvider}s that have been registered by this {@code
+ * RequestStateContext}.
+ */
+ private Set<RequestStateProvider> requestStateProviders = new HashSet<>();
+
+ private RequestStateContext() {}
+
+ /**
+ * Registers a {@link RequestStateProvider}.
+ *
+ * @param requestStateProvider the {@link RequestStateProvider} that should be registered
+ * @return the {@code RequestStateContext} instance for chaining calls
+ */
+ public RequestStateContext addRequestStateProvider(RequestStateProvider requestStateProvider) {
+ if (threadLocalRequestStateProviders.get() == null) {
+ threadLocalRequestStateProviders.set(new HashSet<>());
+ }
+ if (threadLocalRequestStateProviders.get().add(requestStateProvider)) {
+ requestStateProviders.add(requestStateProvider);
+ }
+ return this;
+ }
+
+ /**
+ * Closes this {@code RequestStateContext}.
+ *
+ * <p>Ensures that all {@link RequestStateProvider}s that have been registered by this {@code
+ * RequestStateContext} instance are removed from {@link #threadLocalRequestStateProviders}.
+ *
+ * <p>If no {@link RequestStateProvider}s remain in {@link #threadLocalRequestStateProviders},
+ * {@link #threadLocalRequestStateProviders} is unset.
+ */
+ @Override
+ public void close() {
+ if (threadLocalRequestStateProviders.get() != null) {
+ requestStateProviders.forEach(
+ requestStateProvider ->
+ threadLocalRequestStateProviders.get().remove(requestStateProvider));
+ if (threadLocalRequestStateProviders.get().isEmpty()) {
+ threadLocalRequestStateProviders.remove();
+ }
+ }
+ }
+}
diff --git a/javatests/com/google/gerrit/server/cancellation/RequestStateContextTest.java b/javatests/com/google/gerrit/server/cancellation/RequestStateContextTest.java
new file mode 100644
index 0000000..32903e6
--- /dev/null
+++ b/javatests/com/google/gerrit/server/cancellation/RequestStateContextTest.java
@@ -0,0 +1,106 @@
+// Copyright (C) 2021 The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.google.gerrit.server.cancellation;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import com.google.common.collect.ImmutableSet;
+import org.junit.Test;
+
+public class RequestStateContextTest {
+ @Test
+ public void openContext() {
+ assertNoRequestStateProviders();
+
+ RequestStateProvider requestStateProvider1 = new TestRequestStateProvider();
+ try (RequestStateContext requestStateContext =
+ RequestStateContext.open().addRequestStateProvider(requestStateProvider1)) {
+ RequestStateProvider requestStateProvider2 = new TestRequestStateProvider();
+ requestStateContext.addRequestStateProvider(requestStateProvider2);
+ assertRequestStateProviders(ImmutableSet.of(requestStateProvider1, requestStateProvider2));
+ }
+
+ assertNoRequestStateProviders();
+ }
+
+ @Test
+ public void openNestedContexts() {
+ assertNoRequestStateProviders();
+
+ RequestStateProvider requestStateProvider1 = new TestRequestStateProvider();
+ try (RequestStateContext requestStateContext =
+ RequestStateContext.open().addRequestStateProvider(requestStateProvider1)) {
+ RequestStateProvider requestStateProvider2 = new TestRequestStateProvider();
+ requestStateContext.addRequestStateProvider(requestStateProvider2);
+ assertRequestStateProviders(ImmutableSet.of(requestStateProvider1, requestStateProvider2));
+
+ RequestStateProvider requestStateProvider3 = new TestRequestStateProvider();
+ try (RequestStateContext requestStateContext2 =
+ RequestStateContext.open().addRequestStateProvider(requestStateProvider3)) {
+ RequestStateProvider requestStateProvider4 = new TestRequestStateProvider();
+ requestStateContext2.addRequestStateProvider(requestStateProvider4);
+ assertRequestStateProviders(
+ ImmutableSet.of(
+ requestStateProvider1,
+ requestStateProvider2,
+ requestStateProvider3,
+ requestStateProvider4));
+ }
+
+ assertRequestStateProviders(ImmutableSet.of(requestStateProvider1, requestStateProvider2));
+ }
+
+ assertNoRequestStateProviders();
+ }
+
+ @Test
+ public void openNestedContextsWithSameRequestStateProviders() {
+ assertNoRequestStateProviders();
+
+ RequestStateProvider requestStateProvider1 = new TestRequestStateProvider();
+ try (RequestStateContext requestStateContext =
+ RequestStateContext.open().addRequestStateProvider(requestStateProvider1)) {
+ RequestStateProvider requestStateProvider2 = new TestRequestStateProvider();
+ requestStateContext.addRequestStateProvider(requestStateProvider2);
+ assertRequestStateProviders(ImmutableSet.of(requestStateProvider1, requestStateProvider2));
+
+ try (RequestStateContext requestStateContext2 =
+ RequestStateContext.open().addRequestStateProvider(requestStateProvider1)) {
+ requestStateContext2.addRequestStateProvider(requestStateProvider2);
+
+ assertRequestStateProviders(ImmutableSet.of(requestStateProvider1, requestStateProvider2));
+ }
+
+ assertRequestStateProviders(ImmutableSet.of(requestStateProvider1, requestStateProvider2));
+ }
+
+ assertNoRequestStateProviders();
+ }
+
+ private void assertNoRequestStateProviders() {
+ assertRequestStateProviders(ImmutableSet.of());
+ }
+
+ private void assertRequestStateProviders(
+ ImmutableSet<RequestStateProvider> expectedRequestStateProviders) {
+ assertThat(RequestStateContext.getRequestStateProviders())
+ .containsExactlyElementsIn(expectedRequestStateProviders);
+ }
+
+ private static class TestRequestStateProvider implements RequestStateProvider {
+ @Override
+ public void checkIfCancelled(OnCancelled onCancelled) {}
+ }
+}
