Line data    Source code

       1             : // Copyright (C) 2021 The Android Open Source Project
       2             : //
       3             : // Licensed under the Apache License, Version 2.0 (the "License");
       4             : // you may not use this file except in compliance with the License.
       5             : // You may obtain a copy of the License at
       6             : //
       7             : // http://www.apache.org/licenses/LICENSE-2.0
       8             : //
       9             : // Unless required by applicable law or agreed to in writing, software
      10             : // distributed under the License is distributed on an "AS IS" BASIS,
      11             : // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
      12             : // See the License for the specific language governing permissions and
      13             : // limitations under the License.
      14             : 
      15             : package com.google.gerrit.server.cancellation;
      16             : 
      17             : import com.google.common.annotations.VisibleForTesting;
      18             : import com.google.common.collect.ImmutableSet;
      19             : import java.util.HashSet;
      20             : import java.util.Set;
      21             : 
      22             : /**
      23             :  * Context that allows to register {@link RequestStateProvider}s.
      24             :  *
      25             :  * <p>The registered {@link RequestStateProvider}s are stored in {@link ThreadLocal} so that they
      26             :  * can be accessed during the request execution (via {@link #getRequestStateProviders()}.
      27             :  *
      28             :  * <p>On {@link #close()} the {@link RequestStateProvider}s that have been registered by this {@code
      29             :  * RequestStateContext} instance are removed from {@link ThreadLocal}.
      30             :  *
      31             :  * <p>Nesting {@code RequestStateContext}s is possible.
      32             :  *
      33             :  * <p>Currently there is no logic to automatically copy the {@link RequestStateContext} to
      34             :  * background threads, but implementing this may be considered in the future. This means that by
      35             :  * default we only support cancellation of the main thread, but not of background threads. That's
      36             :  * fine as all significant work is being done in the main thread.
      37             :  *
      38             :  * <p>{@link com.google.gerrit.server.util.RequestContext} is also a context that is available for
      39             :  * the time of the request, but it is not suitable to manage registrations of {@link
      40             :  * RequestStateProvider}s. Hence {@link RequestStateProvider} registrations are managed by a
      41             :  * separate context, which is this class, {@link RequestStateContext}:
      42             :  *
      43             :  * <ul>
      44             :  *   <li>{@link com.google.gerrit.server.util.RequestContext} is an interface that has many
      45             :  *       implementations and hence cannot manage a {@link ThreadLocal} state.
      46             :  *   <li>{@link com.google.gerrit.server.util.RequestContext} is not an {@link AutoCloseable} and
      47             :  *       hence cannot cleanup any {@link ThreadLocal} state on close (turning it into an {@link
      48             :  *       AutoCloseable} would require a large refactoring).
      49             :  *   <li>Despite the name {@link com.google.gerrit.server.util.RequestContext} is not only used for
      50             :  *       requests scopes but also for other scopes that are not a request (e.g. plugin invocations,
      51             :  *       email sending, manual scopes).
      52             :  *   <li>{@link com.google.gerrit.server.util.RequestContext} is not copied to background and should
      53             :  *       not be, but for {@link RequestStateContext} we may consider doing this in the future.
      54             :  * </ul>
      55             :  */
      56             : public class RequestStateContext implements AutoCloseable {
      57             :   /** The {@link RequestStateProvider}s that have been registered for the thread. */
      58         154 :   private static final ThreadLocal<Set<RequestStateProvider>> threadLocalRequestStateProviders =
      59             :       new ThreadLocal<>();
      60             : 
      61             :   /** Whether currently a non-cancellable operation is being performed. */
      62         154 :   private static final ThreadLocal<Boolean> inNonCancellableOperation = new ThreadLocal<>();
      63             : 
      64             :   /**
      65             :    * Aborts the current request by throwing a {@link RequestCancelledException} if any of the
      66             :    * registered {@link RequestStateProvider}s reports the request as cancelled.
      67             :    *
      68             :    * <p>If an atomic operation is currently being performed, request cancellations are ignored and
      69             :    * the request doesn't get aborted.
      70             :    *
      71             :    * @throws RequestCancelledException thrown if the current request is cancelled and should be
      72             :    *     aborted
      73             :    * @see #startNonCancellableOperation()
      74             :    */
      75             :   public static void abortIfCancelled() throws RequestCancelledException {
      76         154 :     if (inNonCancellableOperation.get() != null && inNonCancellableOperation.get()) {
      77             :       // Do not cancel the request while an atomic operation is being performed.
      78         109 :       return;
      79             :     }
      80             : 
      81         154 :     getRequestStateProviders()
      82         154 :         .forEach(
      83             :             requestStateProvider ->
      84         105 :                 requestStateProvider.checkIfCancelled(
      85             :                     (reason, message) -> {
      86           3 :                       throw new RequestCancelledException(reason, message);
      87             :                     }));
      88         154 :   }
      89             : 
      90             :   /**
      91             :    * Starts a non-cancellable operation.
      92             :    *
      93             :    * <p>If the request was cancelled while the non-cancellable operation was running, it gets
      94             :    * aborted on close of the returned {@link AutoCloseable}.
      95             :    *
      96             :    * @return {@link AutoCloseable} that finishes the non-cancellable operation on close.
      97             :    */
      98             :   public static NonCancellableOperationContext startNonCancellableOperation() {
      99         109 :     if (inNonCancellableOperation.get() != null && inNonCancellableOperation.get()) {
     100             :       // atomic operation is already in progress
     101           1 :       return () -> {};
     102             :     }
     103             : 
     104         109 :     inNonCancellableOperation.set(true);
     105         109 :     return () -> {
     106         109 :       inNonCancellableOperation.remove();
     107         109 :       abortIfCancelled();
     108         109 :     };
     109             :   }
     110             : 
     111             :   /** Returns the {@link RequestStateProvider}s that have been registered for the thread. */
     112             :   @VisibleForTesting
     113             :   static ImmutableSet<RequestStateProvider> getRequestStateProviders() {
     114         154 :     if (threadLocalRequestStateProviders.get() == null) {
     115         154 :       return ImmutableSet.of();
     116             :     }
     117         105 :     return ImmutableSet.copyOf(threadLocalRequestStateProviders.get());
     118             :   }
     119             : 
     120             :   /** Opens a {@code RequestStateContext}. */
     121             :   public static RequestStateContext open() {
     122         105 :     return new RequestStateContext();
     123             :   }
     124             : 
     125             :   /**
     126             :    * The {@link RequestStateProvider}s that have been registered by this {@code
     127             :    * RequestStateContext}.
     128             :    */
     129         105 :   private Set<RequestStateProvider> requestStateProviders = new HashSet<>();
     130             : 
     131         105 :   private RequestStateContext() {}
     132             : 
     133             :   /**
     134             :    * Registers a {@link RequestStateProvider}.
     135             :    *
     136             :    * @param requestStateProvider the {@link RequestStateProvider} that should be registered
     137             :    * @return the {@code RequestStateContext} instance for chaining calls
     138             :    */
     139             :   public RequestStateContext addRequestStateProvider(RequestStateProvider requestStateProvider) {
     140         105 :     if (threadLocalRequestStateProviders.get() == null) {
     141         105 :       threadLocalRequestStateProviders.set(new HashSet<>());
     142             :     }
     143         105 :     if (threadLocalRequestStateProviders.get().add(requestStateProvider)) {
     144         105 :       requestStateProviders.add(requestStateProvider);
     145             :     }
     146         105 :     return this;
     147             :   }
     148             : 
     149             :   /**
     150             :    * Closes this {@code RequestStateContext}.
     151             :    *
     152             :    * <p>Ensures that all {@link RequestStateProvider}s that have been registered by this {@code
     153             :    * RequestStateContext} instance are removed from {@link #threadLocalRequestStateProviders}.
     154             :    *
     155             :    * <p>If no {@link RequestStateProvider}s remain in {@link #threadLocalRequestStateProviders},
     156             :    * {@link #threadLocalRequestStateProviders} is unset.
     157             :    */
     158             :   @Override
     159             :   public void close() {
     160         105 :     if (threadLocalRequestStateProviders.get() != null) {
     161         105 :       requestStateProviders.forEach(
     162             :           requestStateProvider ->
     163         105 :               threadLocalRequestStateProviders.get().remove(requestStateProvider));
     164         105 :       if (threadLocalRequestStateProviders.get().isEmpty()) {
     165         105 :         threadLocalRequestStateProviders.remove();
     166             :       }
     167             :     }
     168         105 :   }
     169             : 
     170             :   /**
     171             :    * Context for running a non-cancellable operation.
     172             :    *
     173             :    * <p>While open, the current request cannot be cancelled.
     174             :    */
     175             :   public interface NonCancellableOperationContext extends AutoCloseable {
     176             :     @Override
     177             :     void close();
     178             :   }
     179             : }