| /* |
| * Licensed to the Apache Software Foundation (ASF) under one |
| * or more contributor license agreements. See the NOTICE file |
| * distributed with this work for additional information |
| * regarding copyright ownership. The ASF licenses this file |
| * to you 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 org.apache.cassandra.utils; |
| |
| import java.util.concurrent.CompletableFuture; |
| import java.util.concurrent.ExecutionException; |
| import java.util.concurrent.ExecutorService; |
| import java.util.concurrent.TimeUnit; |
| import java.util.concurrent.TimeoutException; |
| import java.util.concurrent.atomic.AtomicBoolean; |
| import java.util.concurrent.atomic.AtomicReference; |
| import java.util.function.Supplier; |
| |
| /** |
| * Supplier that caches the last computed value until it is reset, forcing every caller of |
| * {@link RecomputingSupplier#get(long, TimeUnit)} to wait until this value is computed if |
| * it was not computed yet. |
| * |
| * Calling {@link RecomputingSupplier#recompute()} won't reset value for the already |
| * waiting consumers, but instead will schedule one recomputation as soon as current one is done. |
| */ |
| public class RecomputingSupplier<T> |
| { |
| private final Supplier<T> supplier; |
| private final AtomicReference<CompletableFuture<T>> cached = new AtomicReference<>(null); |
| private final AtomicBoolean workInProgress = new AtomicBoolean(false); |
| private final ExecutorService executor; |
| |
| public RecomputingSupplier(Supplier<T> supplier, ExecutorService executor) |
| { |
| this.supplier = supplier; |
| this.executor = executor; |
| } |
| |
| public void recompute() |
| { |
| CompletableFuture<T> current = cached.get(); |
| boolean origWip = workInProgress.get(); |
| |
| if (origWip || (current != null && !current.isDone())) |
| { |
| if (cached.get() != current) |
| executor.submit(this::recompute); |
| return; // if work is has not started yet, schedule task for the future |
| } |
| |
| assert current == null || current.isDone(); |
| |
| // The work is not in progress, and current future is done. Try to submit a new task. |
| CompletableFuture<T> lazyValue = new CompletableFuture<>(); |
| if (cached.compareAndSet(current, lazyValue)) |
| executor.submit(() -> doWork(lazyValue)); |
| else |
| executor.submit(this::recompute); // Lost CAS, resubmit |
| } |
| |
| private void doWork(CompletableFuture<T> lazyValue) |
| { |
| T value = null; |
| Throwable err = null; |
| try |
| { |
| sanityCheck(workInProgress.compareAndSet(false, true)); |
| value = supplier.get(); |
| } |
| catch (Throwable t) |
| { |
| err = t; |
| } |
| finally |
| { |
| sanityCheck(workInProgress.compareAndSet(true, false)); |
| } |
| |
| if (err == null) |
| lazyValue.complete(value); |
| else |
| lazyValue.completeExceptionally(err); |
| } |
| |
| private static void sanityCheck(boolean check) |
| { |
| assert check : "At most one task should be executing using this executor"; |
| } |
| |
| public T get(long timeout, TimeUnit timeUnit) throws InterruptedException, ExecutionException, TimeoutException |
| { |
| CompletableFuture<T> lazyValue = cached.get(); |
| |
| // recompute was never called yet, return null. |
| if (lazyValue == null) |
| return null; |
| |
| return lazyValue.get(timeout, timeUnit); |
| } |
| |
| public String toString() |
| { |
| return "RecomputingSupplier{" + |
| "supplier=" + supplier + |
| ", cached=" + cached + |
| ", workInProgress=" + workInProgress + |
| ", executor=" + executor + |
| '}'; |
| } |
| } |