RATIS-2497. Pass server to the dummy watch request in OrderedAsync

Author: szetszwo

ratis-client/src/main/java/org/apache/ratis/client/impl/AsyncImpl.java

@@ -40,7 +40,7 @@ class AsyncImpl implements AsyncRpcApi {
   CompletableFuture<RaftClientReply> send(
       RaftClientRequest.Type type, Message message, RaftPeerId server) {
     return TraceClient.asyncSend(
-        () -> client.getOrderedAsync().send(type, message, server), type, server);
+        () -> client.getOrderedAsync(server).send(type, message, server), type, server);
   }
 
   @Override

ratis-client/src/main/java/org/apache/ratis/client/impl/OrderedAsync.java

@@ -116,12 +116,12 @@ public String toString() {
     }
   }
 
-  static OrderedAsync newInstance(RaftClientImpl client, RaftProperties properties) {
+  static OrderedAsync newInstance(RaftClientImpl client, RaftPeerId server, RaftProperties properties) {
     final OrderedAsync ordered = new OrderedAsync(client, properties);
     // send a dummy watch request to establish the connection
     // TODO: this is a work around, it is better to fix the underlying RPC implementation
     if (RaftClientConfigKeys.Async.Experimental.sendDummyRequest(properties)) {
-      ordered.send(RaftClientRequest.watchRequestType(), null, null);
+      ordered.send(RaftClientRequest.watchRequestType(), null, server);
     }
     return ordered;
   }

ratis-client/src/main/java/org/apache/ratis/client/impl/RaftClientImpl.java

@@ -47,6 +47,7 @@
 import org.apache.ratis.util.CollectionUtils;
 import org.apache.ratis.util.IOUtils;
 import org.apache.ratis.util.JavaUtils;
+import org.apache.ratis.util.MemoizedFunction;
 import org.apache.ratis.util.MemoizedSupplier;
 import org.apache.ratis.util.Preconditions;
 import org.apache.ratis.util.TimeDuration;
@@ -180,7 +181,7 @@ private synchronized Set<Long> getAndReset() {
 
   private final TimeoutExecutor scheduler = TimeoutExecutor.getInstance();
 
-  private final Supplier<OrderedAsync> orderedAsync;
+  private final MemoizedFunction<RaftPeerId, OrderedAsync> orderedAsync;
   private final Supplier<AsyncImpl> asyncApi;
   private final Supplier<BlockingImpl> blockingApi;
   private final Supplier<MessageStreamImpl> messageStreamApi;
@@ -209,7 +210,7 @@ private synchronized Set<Long> getAndReset() {
     clientRpc.addRaftPeers(group.getPeers());
     this.clientRpc = clientRpc;
 
-    this.orderedAsync = JavaUtils.memoize(() -> OrderedAsync.newInstance(this, properties));
+    this.orderedAsync = MemoizedFunction.valueOf(server -> OrderedAsync.newInstance(this, server, properties));
     this.messageStreamApi = JavaUtils.memoize(() -> MessageStreamImpl.newInstance(this, properties));
     this.asyncApi = JavaUtils.memoize(() -> new AsyncImpl(this));
     this.blockingApi = JavaUtils.memoize(() -> new BlockingImpl(this));
@@ -277,8 +278,8 @@ TimeoutExecutor getScheduler() {
     return scheduler;
   }
 
-  OrderedAsync getOrderedAsync() {
-    return orderedAsync.get();
+  OrderedAsync getOrderedAsync(RaftPeerId server) {
+    return orderedAsync.apply(server);
   }
 
   RaftClientRequest newRaftClientRequest(

ratis-common/src/main/java/org/apache/ratis/util/MemoizedBase.java

@@ -0,0 +1,73 @@
+/*
+ * 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.ratis.util;
+
+import org.apache.ratis.util.function.CheckedSupplier;
+
+import java.util.Objects;
+
+/**
+ * This is the base class for the memoized subclass such as
+ * {@link MemoizedSupplier}, {@link MemoizedFunction}, {@link MemoizedCheckedSupplier}, etc,
+ * The subclasses provide its own method to retrieve the value,
+ * such as {@link MemoizedSupplier#get()} and {@link MemoizedFunction#apply(Object)}.
+ * The subclass method returns a value by invoking its initializer once at the first call
+ * and then keeps returning the same value for the subsequent calls.
+ * <p>
+ * All the subclasses are thread safe.
+ *
+ * @param <RETURN> The value type.
+ * @param <THROW> The throwable type of the initializer.
+ */
+abstract class MemoizedBase<RETURN, THROW extends Throwable> {
+  @SuppressWarnings({"squid:S3077"}) // Suppress volatile for generic type
+  private volatile RETURN value = null;
+
+  final RETURN init(CheckedSupplier<RETURN, THROW> initializer) throws THROW {
+    final RETURN initialized = value;
+    if (initialized != null) {
+      return initialized;
+    }
+
+    synchronized (this) {
+      if (value == null) {
+        value = initializer.get();
+        Objects.requireNonNull(value, "initializer.get() returns null");
+      }
+      return value;
+    }
+  }
+
+  /** @return is the object initialized? */
+  public final boolean isInitialized() {
+    return value != null;
+  }
+
+  /**
+   * @return the value, which must be already initialized.
+   * @throws NullPointerException if the value is uninitialized.
+   */
+  public RETURN getInitializedValue() {
+    return Objects.requireNonNull(value, "Uninitialized: value == null");
+  }
+
+  @Override
+  public String toString() {
+    return value != null ? "Memoized:" + value : "Uninitialized";
+  }
+}

ratis-common/src/main/java/org/apache/ratis/util/MemoizedCheckedSupplier.java

@@ -1,4 +1,4 @@
-/**
+/*
  * 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
@@ -25,13 +25,14 @@
  * A memoized supplier is a {@link CheckedSupplier}
  * which gets a value by invoking its initializer once.
  * and then keeps returning the same value as its supplied results.
- *
+ * <p>
  * This class is thread safe.
  *
  * @param <RETURN> The return type of the supplier.
  * @param <THROW> The throwable type of the supplier.
  */
 public final class MemoizedCheckedSupplier<RETURN, THROW extends Throwable>
+    extends MemoizedBase<RETURN, THROW>
     implements CheckedSupplier<RETURN, THROW> {
   /**
    * @param supplier to supply at most one non-null value.
@@ -45,9 +46,6 @@ public static <RETURN, THROW extends Throwable> MemoizedCheckedSupplier<RETURN,
 
   private final CheckedSupplier<RETURN, THROW> initializer;
 
-  @SuppressWarnings({"squid:S3077"}) // Suppress volatile for generic type
-  private volatile RETURN value = null;
-
   /**
    * Create a memoized supplier.
    * @param initializer to supply at most one non-null value.
@@ -60,33 +58,14 @@ private MemoizedCheckedSupplier(CheckedSupplier<RETURN, THROW> initializer) {
   /** @return the lazily initialized object. */
   @Override
   public RETURN get() throws THROW {
-    RETURN v = value;
-    if (v == null) {
-      synchronized (this) {
-        v = value;
-        if (v == null) {
-          v = value = Objects.requireNonNull(initializer.get(), "initializer.get() returns null");
-        }
-      }
-    }
-    return v;
+    return init(initializer);
   }
 
   /**
    * @return the already initialized object.
    * @throws NullPointerException if the object is uninitialized.
    */
   public RETURN getUnchecked() {
-    return Objects.requireNonNull(value, "value == null");
-  }
-
-  /** @return is the object initialized? */
-  public boolean isInitialized() {
-    return value != null;
-  }
-
-  @Override
-  public String toString() {
-    return isInitialized()? "Memoized:" + value: "UNINITIALIZED";
+    return getInitializedValue();
   }
 }

ratis-common/src/main/java/org/apache/ratis/util/MemoizedFunction.java

@@ -0,0 +1,69 @@
+/*
+ * 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.ratis.util;
+
+import java.util.Objects;
+import java.util.function.Function;
+
+/**
+ * A memoized function is a {@link Function}
+ * which returns a value by invoking its initializer once
+ * and then keeps returning the same value as its result.
+ * <p>
+ * This class is similar to {@link MemoizedSupplier} except that
+ * the initializer takes a parameter.
+ * <p>
+ * This class is thread safe.
+ *
+ * @param <RETURN> The function result type.
+ */
+public final class MemoizedFunction<PARAMETER, RETURN>
+    extends MemoizedBase<RETURN, RuntimeException>
+    implements Function<PARAMETER, RETURN> {
+  /**
+   * @param function to supply at most one non-null value.
+   * @return a {@link MemoizedFunction} with the given function.
+   */
+  public static <P, R> MemoizedFunction<P, R> valueOf(Function<P, R> function) {
+    return function instanceof MemoizedFunction ?
+        (MemoizedFunction<P, R>) function : new MemoizedFunction<>(function);
+  }
+
+  private final Function<PARAMETER, RETURN> initializer;
+
+  /**
+   * Create a memoized function.
+   * @param initializer to supply at most one non-null value.
+   */
+  private MemoizedFunction(Function<PARAMETER, RETURN> initializer) {
+    Objects.requireNonNull(initializer, "initializer == null");
+    this.initializer = initializer;
+  }
+
+  /**
+   * @param parameter for passing to the initializer.
+   *                  Since the returned function is memoized, the parameter is only used at the first call.
+   *                  The parameter in the subsequent calls is ignored.
+   *
+   * @return the lazily initialized object.
+   */
+  @Override
+  public RETURN apply(PARAMETER parameter) {
+    return init(() -> initializer.apply(parameter));
+  }
+}

ratis-common/src/main/java/org/apache/ratis/util/MemoizedSupplier.java

@@ -1,4 +1,4 @@
-/**
+/*
  * 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
@@ -24,12 +24,14 @@
  * A memoized supplier is a {@link Supplier}
  * which gets a value by invoking its initializer once
  * and then keeps returning the same value as its supplied results.
- *
+ * <p>
  * This class is thread safe.
  *
  * @param <T> The supplier result type.
  */
-public final class MemoizedSupplier<T> implements Supplier<T> {
+public final class MemoizedSupplier<T>
+    extends MemoizedBase<T, RuntimeException>
+    implements Supplier<T> {
   /**
    * @param supplier to supply at most one non-null value.
    * @return a {@link MemoizedSupplier} with the given supplier.
@@ -40,8 +42,6 @@ public static <T> MemoizedSupplier<T> valueOf(Supplier<T> supplier) {
   }
 
   private final Supplier<T> initializer;
-  @SuppressWarnings({"squid:S3077"}) // Suppress volatile for generic type
-  private volatile T value = null;
 
   /**
    * Create a memoized supplier.
@@ -55,26 +55,6 @@ private MemoizedSupplier(Supplier<T> initializer) {
   /** @return the lazily initialized object. */
   @Override
   public T get() {
-    T v = value;
-    if (v == null) {
-      synchronized (this) {
-        v = value;
-        if (v == null) {
-          v = value = Objects.requireNonNull(initializer.get(),
-              "initializer.get() returns null");
-        }
-      }
-    }
-    return v;
-  }
-
-  /** @return is the object initialized? */
-  public boolean isInitialized() {
-    return value != null;
-  }
-
-  @Override
-  public String toString() {
-    return isInitialized()? "Memoized:" + get(): "UNINITIALIZED";
+    return init(initializer::get);
   }
 }

ratis-server/src/test/java/org/apache/ratis/client/impl/RaftClientTestUtil.java

@@ -29,7 +29,8 @@
 /** Interface for testing raft client. */
 public interface RaftClientTestUtil {
   static void assertAsyncRequestSemaphore(RaftClient client, int expectedAvailablePermits, int expectedQueueLength) {
-    ((RaftClientImpl) client).getOrderedAsync().assertRequestSemaphore(expectedAvailablePermits, expectedQueueLength);
+    ((RaftClientImpl) client).getOrderedAsync(null)
+        .assertRequestSemaphore(expectedAvailablePermits, expectedQueueLength);
   }
 
   static ClientInvocationId getClientInvocationId(RaftClient client) {