[FLINK-19130] [core] Expose metrics for backpressure

statefun-flink/statefun-flink-core/src/main/java/org/apache/flink/statefun/flink/core/backpressure/BackPressureValve.java

@@ -53,4 +53,12 @@ public interface BackPressureValve {
    * @param address the address
    */
   void blockAddress(Address address);
+
+  /**
+   * Checks whether a given address was previously blocked with {@link #blockAddress(Address)}.
+   *
+   * @param address the address to check
+   * @return boolean indicating whether or not the address was blocked.
+   */
+  boolean isAddressBlocked(Address address);
 }

statefun-flink/statefun-flink-core/src/main/java/org/apache/flink/statefun/flink/core/backpressure/AsyncWaiter.java → statefun-flink/statefun-flink-core/src/main/java/org/apache/flink/statefun/flink/core/backpressure/InternalContext.java

@@ -19,9 +19,11 @@
 package org.apache.flink.statefun.flink.core.backpressure;
 
 import org.apache.flink.annotation.Internal;
+import org.apache.flink.statefun.flink.core.metrics.FunctionTypeMetrics;
+import org.apache.flink.statefun.sdk.Context;
 
 @Internal
-public interface AsyncWaiter {
+public interface InternalContext extends Context {
 
   /**
    * Signals the runtime to stop invoking the currently executing function with new input until at
@@ -36,4 +38,11 @@ public interface AsyncWaiter {
    * every async operation registered per each address.
    */
   void awaitAsyncOperationComplete();
+
+  /**
+   * Returns the metrics handle for the current invoked function's type.
+   *
+   * @return the metrics handle for the current invoked function's type.
+   */
+  FunctionTypeMetrics functionTypeMetrics();
 }

statefun-flink/statefun-flink-core/src/main/java/org/apache/flink/statefun/flink/core/backpressure/ThresholdBackPressureValve.java

@@ -43,7 +43,7 @@ public final class ThresholdBackPressureValve implements BackPressureValve {
 
   /**
    * a set of address that had explicitly requested to stop processing any new inputs (via {@link
-   * AsyncWaiter#awaitAsyncOperationComplete()}. Note that this is a set implemented on top of a
+   * InternalContext#awaitAsyncOperationComplete()}. Note that this is a set implemented on top of a
    * map, and the value (Boolean) has no meaning.
    */
   private final ObjectOpenHashMap<Address, Boolean> blockedAddressSet =
@@ -88,6 +88,12 @@ public void notifyAsyncOperationCompleted(Address owningAddress) {
     blockedAddressSet.remove(owningAddress);
   }
 
+  /** {@inheritDoc} */
+  @Override
+  public boolean isAddressBlocked(Address address) {
+    return blockedAddressSet.containsKey(address);
+  }
+
   private boolean totalPendingAsyncOperationsAtCapacity() {
     return maximumPendingAsynchronousOperations > 0
         && pendingAsynchronousOperationsCount >= maximumPendingAsynchronousOperations;

statefun-flink/statefun-flink-core/src/main/java/org/apache/flink/statefun/flink/core/di/ObjectContainer.java

@@ -51,6 +51,14 @@ public <T> void add(String label, Class<? super T> type, Class<?> actual) {
     factories.put(new Key(type, label), () -> createReflectively(actual));
   }
 
+  public <T, ET> void addAlias(
+      String newLabel,
+      Class<? super T> newType,
+      String existingLabel,
+      Class<? super ET> existingType) {
+    factories.put(new Key(newType, newLabel), () -> get(existingType, existingLabel));
+  }
+
   public <T> void add(String label, Lazy<T> lazyValue) {
     factories.put(new Key(Lazy.class, label), () -> lazyValue.withContainer(this));
   }

statefun-flink/statefun-flink-core/src/main/java/org/apache/flink/statefun/flink/core/functions/AsyncSink.java

@@ -27,6 +27,9 @@
 import org.apache.flink.statefun.flink.core.di.Label;
 import org.apache.flink.statefun.flink.core.di.Lazy;
 import org.apache.flink.statefun.flink.core.message.Message;
+import org.apache.flink.statefun.flink.core.metrics.FunctionDispatcherMetrics;
+import org.apache.flink.statefun.flink.core.metrics.FunctionTypeMetrics;
+import org.apache.flink.statefun.flink.core.metrics.FunctionTypeMetricsRepository;
 import org.apache.flink.statefun.flink.core.queue.Locks;
 import org.apache.flink.statefun.flink.core.queue.MpscQueue;
 import org.apache.flink.statefun.sdk.Address;
@@ -36,6 +39,8 @@ final class AsyncSink {
   private final Lazy<Reductions> reductions;
   private final Executor operatorMailbox;
   private final BackPressureValve backPressureValve;
+  private final FunctionTypeMetricsRepository metricsRepository;
+  private final FunctionDispatcherMetrics dispatcherMetrics;
 
   private final MpscQueue<Message> completed = new MpscQueue<>(32768, Locks.jdkReentrantLock());
 
@@ -44,23 +49,30 @@ final class AsyncSink {
       PendingAsyncOperations pendingAsyncOperations,
       @Label("mailbox-executor") Executor operatorMailbox,
       @Label("reductions") Lazy<Reductions> reductions,
-      @Label("backpressure-valve") BackPressureValve backPressureValve) {
+      @Label("backpressure-valve") BackPressureValve backPressureValve,
+      @Label("function-metrics-repository") FunctionTypeMetricsRepository metricsRepository,
+      @Label("function-dispatcher-metrics") FunctionDispatcherMetrics dispatcherMetrics) {
     this.pendingAsyncOperations = Objects.requireNonNull(pendingAsyncOperations);
     this.reductions = Objects.requireNonNull(reductions);
     this.operatorMailbox = Objects.requireNonNull(operatorMailbox);
     this.backPressureValve = Objects.requireNonNull(backPressureValve);
+    this.metricsRepository = Objects.requireNonNull(metricsRepository);
+    this.dispatcherMetrics = Objects.requireNonNull(dispatcherMetrics);
   }
 
-  <T> void accept(Message metadata, CompletableFuture<T> future) {
+  <T> void accept(Address sourceAddress, Message metadata, CompletableFuture<T> future) {
     final long futureId = ThreadLocalRandom.current().nextLong(); // TODO: is this is good enough?
     // we keep the message in state (associated with futureId) until either:
     // 1. the future successfully completes and the message is processed. The state would be
     // cleared by the AsyncMessageDecorator after a successful application.
     // 2. after recovery, we clear that state by notifying the owning function that we don't know
     // what happened
     // with that particular async operation.
-    pendingAsyncOperations.add(metadata.source(), futureId, metadata);
+    pendingAsyncOperations.add(sourceAddress, futureId, metadata);
     backPressureValve.notifyAsyncOperationRegistered();
+
+    metricsRepository.getMetrics(sourceAddress.type()).asyncOperationRegistered();
+    dispatcherMetrics.asyncOperationRegistered();
     future.whenComplete((result, throwable) -> enqueue(metadata, futureId, result, throwable));
   }
 
@@ -72,6 +84,7 @@ <T> void accept(Message metadata, CompletableFuture<T> future) {
    */
   void blockAddress(Address address) {
     backPressureValve.blockAddress(address);
+    metricsRepository.getMetrics(address.type()).blockedAddress();
   }
 
   private <T> void enqueue(Message message, long futureId, T result, Throwable throwable) {
@@ -90,7 +103,17 @@ private void drainOnOperatorThread() {
     Reductions reductions = this.reductions.get();
     Message message;
     while ((message = batchOfCompletedFutures.poll()) != null) {
-      backPressureValve.notifyAsyncOperationCompleted(message.target());
+      Address target = message.target();
+      FunctionTypeMetrics functionMetrics = metricsRepository.getMetrics(target.type());
+
+      // must check whether address was blocked BEFORE notifying completion
+      if (backPressureValve.isAddressBlocked(target)) {
+        functionMetrics.unblockedAddress();
+      }
+      backPressureValve.notifyAsyncOperationCompleted(target);
+
+      functionMetrics.asyncOperationCompleted();
+      dispatcherMetrics.asyncOperationCompleted();
       reductions.enqueue(message);
     }
     reductions.processEnvelopes();

statefun-flink/statefun-flink-core/src/main/java/org/apache/flink/statefun/flink/core/functions/Reductions.java

@@ -32,8 +32,11 @@
 import org.apache.flink.statefun.flink.core.di.ObjectContainer;
 import org.apache.flink.statefun.flink.core.message.Message;
 import org.apache.flink.statefun.flink.core.message.MessageFactory;
-import org.apache.flink.statefun.flink.core.metrics.FlinkMetricsFactory;
-import org.apache.flink.statefun.flink.core.metrics.MetricsFactory;
+import org.apache.flink.statefun.flink.core.metrics.FlinkFuncionTypeMetricsFactory;
+import org.apache.flink.statefun.flink.core.metrics.FlinkFunctionDispatcherMetrics;
+import org.apache.flink.statefun.flink.core.metrics.FuncionTypeMetricsFactory;
+import org.apache.flink.statefun.flink.core.metrics.FunctionDispatcherMetrics;
+import org.apache.flink.statefun.flink.core.metrics.FunctionTypeMetricsRepository;
 import org.apache.flink.statefun.flink.core.state.FlinkState;
 import org.apache.flink.statefun.flink.core.state.State;
 import org.apache.flink.statefun.flink.core.types.DynamicallyRegisteredTypes;
@@ -71,6 +74,11 @@ static Reductions create(
     container.add("function-providers", Map.class, statefulFunctionsUniverse.functions());
     container.add(
         "function-repository", FunctionRepository.class, StatefulFunctionRepository.class);
+    container.addAlias(
+        "function-metrics-repository",
+        FunctionTypeMetricsRepository.class,
+        "function-repository",
+        FunctionRepository.class);
 
     // for FlinkState
     container.add("runtime-context", RuntimeContext.class, context);
@@ -95,7 +103,14 @@ static Reductions create(
     container.add("function-loader", FunctionLoader.class, PredefinedFunctionLoader.class);
     container.add(Reductions.class);
     container.add(LocalFunctionGroup.class);
-    container.add("metrics-factory", MetricsFactory.class, new FlinkMetricsFactory(metricGroup));
+    container.add(
+        "function-metrics-factory",
+        FuncionTypeMetricsFactory.class,
+        new FlinkFuncionTypeMetricsFactory(metricGroup));
+    container.add(
+        "function-dispatcher-metrics",
+        FunctionDispatcherMetrics.class,
+        new FlinkFunctionDispatcherMetrics(metricGroup));
 
     // for delayed messages
     container.add(

statefun-flink/statefun-flink-core/src/main/java/org/apache/flink/statefun/flink/core/functions/ReusableContext.java

@@ -20,16 +20,17 @@
 import java.time.Duration;
 import java.util.Objects;
 import java.util.concurrent.CompletableFuture;
-import org.apache.flink.statefun.flink.core.backpressure.AsyncWaiter;
+import org.apache.flink.statefun.flink.core.backpressure.InternalContext;
 import org.apache.flink.statefun.flink.core.di.Inject;
 import org.apache.flink.statefun.flink.core.di.Label;
 import org.apache.flink.statefun.flink.core.message.Message;
 import org.apache.flink.statefun.flink.core.message.MessageFactory;
+import org.apache.flink.statefun.flink.core.metrics.FunctionTypeMetrics;
 import org.apache.flink.statefun.flink.core.state.State;
 import org.apache.flink.statefun.sdk.Address;
 import org.apache.flink.statefun.sdk.io.EgressIdentifier;
 
-final class ReusableContext implements ApplyingContext, AsyncWaiter {
+final class ReusableContext implements ApplyingContext, InternalContext {
   private final Partition thisPartition;
   private final LocalSink localSink;
   private final RemoteSink remoteSink;
@@ -113,14 +114,19 @@ public <M, T> void registerAsyncOperation(M metadata, CompletableFuture<T> futur
     Objects.requireNonNull(future);
 
     Message message = messageFactory.from(self(), self(), metadata);
-    asyncSink.accept(message, future);
+    asyncSink.accept(self(), message, future);
   }
 
   @Override
   public void awaitAsyncOperationComplete() {
     asyncSink.blockAddress(self());
   }
 
+  @Override
+  public FunctionTypeMetrics functionTypeMetrics() {
+    return function.metrics();
+  }
+
   @Override
   public Address caller() {
     return in.source();

statefun-flink/statefun-flink-core/src/main/java/org/apache/flink/statefun/flink/core/functions/StatefulFunctionRepository.java

@@ -23,29 +23,31 @@
 import org.apache.flink.statefun.flink.core.di.Inject;
 import org.apache.flink.statefun.flink.core.di.Label;
 import org.apache.flink.statefun.flink.core.message.MessageFactory;
+import org.apache.flink.statefun.flink.core.metrics.FuncionTypeMetricsFactory;
 import org.apache.flink.statefun.flink.core.metrics.FunctionTypeMetrics;
-import org.apache.flink.statefun.flink.core.metrics.MetricsFactory;
+import org.apache.flink.statefun.flink.core.metrics.FunctionTypeMetricsRepository;
 import org.apache.flink.statefun.flink.core.state.FlinkStateBinder;
 import org.apache.flink.statefun.flink.core.state.PersistedStates;
 import org.apache.flink.statefun.flink.core.state.State;
 import org.apache.flink.statefun.sdk.FunctionType;
 
-final class StatefulFunctionRepository implements FunctionRepository {
+final class StatefulFunctionRepository
+    implements FunctionRepository, FunctionTypeMetricsRepository {
   private final ObjectOpenHashMap<FunctionType, StatefulFunction> instances;
   private final State flinkState;
   private final FunctionLoader functionLoader;
-  private final MetricsFactory metricsFactory;
+  private final FuncionTypeMetricsFactory metricsFactory;
   private final MessageFactory messageFactory;
 
   @Inject
   StatefulFunctionRepository(
       @Label("function-loader") FunctionLoader functionLoader,
-      @Label("metrics-factory") MetricsFactory metricsFactory,
+      @Label("function-metrics-factory") FuncionTypeMetricsFactory functionMetricsFactory,
       @Label("state") State state,
       MessageFactory messageFactory) {
     this.instances = new ObjectOpenHashMap<>();
     this.functionLoader = Objects.requireNonNull(functionLoader);
-    this.metricsFactory = Objects.requireNonNull(metricsFactory);
+    this.metricsFactory = Objects.requireNonNull(functionMetricsFactory);
     this.flinkState = Objects.requireNonNull(state);
     this.messageFactory = Objects.requireNonNull(messageFactory);
   }
@@ -59,6 +61,11 @@ public LiveFunction get(FunctionType type) {
     return function;
   }
 
+  @Override
+  public FunctionTypeMetrics getMetrics(FunctionType functionType) {
+    return get(functionType).metrics();
+  }
+
   private StatefulFunction load(FunctionType functionType) {
     org.apache.flink.statefun.sdk.StatefulFunction statefulFunction =
         functionLoader.load(functionType);

statefun-flink/statefun-flink-core/src/main/java/org/apache/flink/statefun/flink/core/metrics/FlinkMetricsFactory.java → statefun-flink/statefun-flink-core/src/main/java/org/apache/flink/statefun/flink/core/metrics/FlinkFuncionTypeMetricsFactory.java

@@ -21,11 +21,11 @@
 import org.apache.flink.metrics.MetricGroup;
 import org.apache.flink.statefun.sdk.FunctionType;
 
-public class FlinkMetricsFactory implements MetricsFactory {
+public class FlinkFuncionTypeMetricsFactory implements FuncionTypeMetricsFactory {
 
   private final MetricGroup metricGroup;
 
-  public FlinkMetricsFactory(MetricGroup metricGroup) {
+  public FlinkFuncionTypeMetricsFactory(MetricGroup metricGroup) {
     this.metricGroup = Objects.requireNonNull(metricGroup);
   }
 

statefun-flink/statefun-flink-core/src/main/java/org/apache/flink/statefun/flink/core/metrics/FlinkFunctionDispatcherMetrics.java

@@ -0,0 +1,43 @@
+/*
+ * 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.flink.statefun.flink.core.metrics;
+
+import java.util.Objects;
+import org.apache.flink.metrics.Counter;
+import org.apache.flink.metrics.MetricGroup;
+
+public class FlinkFunctionDispatcherMetrics implements FunctionDispatcherMetrics {
+  private final Counter inflightAsyncOperations;
+
+  public FlinkFunctionDispatcherMetrics(MetricGroup operatorGroup) {
+    Objects.requireNonNull(operatorGroup, "operatorGroup");
+
+    this.inflightAsyncOperations = operatorGroup.counter("inflight-async-ops");
+  }
+
+  @Override
+  public void asyncOperationRegistered() {
+    inflightAsyncOperations.inc();
+  }
+
+  @Override
+  public void asyncOperationCompleted() {
+    inflightAsyncOperations.dec();
+  }
+}