SAMZA-1219; Add metrics for operator message received and execution times

Author: Prateek Maheshwari <pmaheshw@linkedin.com>

Reviewers: Jagadish <jagadish@apache.org>

Closes apache#142 from prateekm/operator-metrics

Author: prateekm

samza-core/src/main/java/org/apache/samza/operators/impl/OperatorImpl.java

@@ -18,9 +18,19 @@
  */
 package org.apache.samza.operators.impl;
 
+import org.apache.samza.config.Config;
+import org.apache.samza.config.MetricsConfig;
+import org.apache.samza.metrics.Counter;
+import org.apache.samza.metrics.MetricsRegistry;
+import org.apache.samza.metrics.Timer;
+import org.apache.samza.operators.spec.OperatorSpec;
 import org.apache.samza.task.MessageCollector;
+import org.apache.samza.task.TaskContext;
 import org.apache.samza.task.TaskCoordinator;
+import org.apache.samza.util.HighResolutionClock;
 
+import java.util.Collection;
+import java.util.Collections;
 import java.util.HashSet;
 import java.util.Set;
 
@@ -29,60 +39,140 @@
  * Abstract base class for all stream operator implementations.
  */
 public abstract class OperatorImpl<M, RM> {
+  private static final String METRICS_GROUP = OperatorImpl.class.getName();
 
-  private final Set<OperatorImpl<RM, ?>> nextOperators = new HashSet<>();
+  private boolean initialized;
+  private Set<OperatorImpl<RM, ?>> registeredOperators;
+  private HighResolutionClock highResClock;
+  private Counter numMessage;
+  private Timer handleMessageNs;
+  private Timer handleTimerNs;
 
   /**
-   * Register the next operator in the chain that this operator should propagate its output to.
-   * @param nextOperator  the next operator in the chain.
+   * Initialize this {@link OperatorImpl} and its user-defined functions.
+   *
+   * @param config  the {@link Config} for the task
+   * @param context  the {@link TaskContext} for the task
+   */
+  public final void init(Config config, TaskContext context) {
+    String opName = getOperatorSpec().getOpName();
+
+    if (initialized) {
+      throw new IllegalStateException(String.format("Attempted to initialize Operator %s more than once.", opName));
+    }
+
+    this.highResClock = createHighResClock(config);
+    registeredOperators = new HashSet<>();
+    MetricsRegistry metricsRegistry = context.getMetricsRegistry();
+    this.numMessage = metricsRegistry.newCounter(METRICS_GROUP, opName + "-messages");
+    this.handleMessageNs = metricsRegistry.newTimer(METRICS_GROUP, opName + "-handle-message-ns");
+    this.handleTimerNs = metricsRegistry.newTimer(METRICS_GROUP, opName + "-handle-timer-ns");
+
+    handleInit(config, context);
+
+    initialized = true;
+  }
+
+  /**
+   * Initialize this {@link OperatorImpl} and its user-defined functions.
+   *
+   * @param config  the {@link Config} for the task
+   * @param context  the {@link TaskContext} for the task
+   */
+  protected abstract void handleInit(Config config, TaskContext context);
+
+  /**
+   * Register an operator that this operator should propagate its results to.
+   *
+   * @param nextOperator  the next operator to propagate results to
    */
   void registerNextOperator(OperatorImpl<RM, ?> nextOperator) {
-    nextOperators.add(nextOperator);
+    if (!initialized) {
+      throw new IllegalStateException(
+          String.format("Attempted to register next operator before initializing operator %s.",
+              getOperatorSpec().getOpName()));
+    }
+    this.registeredOperators.add(nextOperator);
   }
 
   /**
-   * Perform the transformation required for this operator and call the downstream operators.
+   * Handle the incoming {@code message} for this {@link OperatorImpl} and propagate results to registered operators.
+   * <p>
+   * Delegates to {@link #handleMessage(Object, MessageCollector, TaskCoordinator)} for handling the message.
    *
-   * Must call {@link #propagateResult} to propagate the output to registered downstream operators correctly.
+   * @param message  the input message
+   * @param collector  the {@link MessageCollector} for this message
+   * @param coordinator  the {@link TaskCoordinator} for this message
+   */
+  public final void onMessage(M message, MessageCollector collector, TaskCoordinator coordinator) {
+    this.numMessage.inc();
+    long startNs = this.highResClock.nanoTime();
+    Collection<RM> results = handleMessage(message, collector, coordinator);
+    long endNs = this.highResClock.nanoTime();
+    this.handleMessageNs.update(endNs - startNs);
+
+    results.forEach(rm ->
+        this.registeredOperators.forEach(op ->
+            op.onMessage(rm, collector, coordinator)));
+  }
+
+  /**
+   * Handle the incoming {@code message} and return the results to be propagated to registered operators.
    *
    * @param message  the input message
    * @param collector  the {@link MessageCollector} in the context
    * @param coordinator  the {@link TaskCoordinator} in the context
+   * @return  results of the transformation
    */
-  public abstract void onNext(M message, MessageCollector collector, TaskCoordinator coordinator);
+  protected abstract Collection<RM> handleMessage(M message, MessageCollector collector,
+      TaskCoordinator coordinator);
 
   /**
-   * Invoked at every tick. This method delegates to {@link #onTimer(MessageCollector, TaskCoordinator)}
+   * Handle timer ticks for this {@link OperatorImpl} and propagate the results and timer tick to registered operators.
+   * <p>
+   * Delegates to {@link #handleTimer(MessageCollector, TaskCoordinator)} for handling the timer tick.
    *
    * @param collector  the {@link MessageCollector} in the context
    * @param coordinator  the {@link TaskCoordinator} in the context
    */
-  public final void onTick(MessageCollector collector, TaskCoordinator coordinator) {
-    onTimer(collector, coordinator);
-    nextOperators.forEach(sub -> sub.onTick(collector, coordinator));
+  public final void onTimer(MessageCollector collector, TaskCoordinator coordinator) {
+    long startNs = this.highResClock.nanoTime();
+    Collection<RM> results = handleTimer(collector, coordinator);
+    long endNs = this.highResClock.nanoTime();
+    this.handleTimerNs.update(endNs - startNs);
+
+    results.forEach(rm ->
+        this.registeredOperators.forEach(op ->
+            op.onMessage(rm, collector, coordinator)));
+    this.registeredOperators.forEach(op ->
+        op.onTimer(collector, coordinator));
   }
 
   /**
-   * Invoked at every tick. Implementations must call {@link #propagateResult} to propagate any generated output
-   * to registered downstream operators.
+   * Handle the the timer tick for this operator and return the results to be propagated to registered operators.
+   * <p>
+   * Defaults to a no-op implementation.
    *
    * @param collector  the {@link MessageCollector} in the context
    * @param coordinator  the {@link TaskCoordinator} in the context
+   * @return  results of the timed operation
    */
-  public void onTimer(MessageCollector collector, TaskCoordinator coordinator) {
+  protected Collection<RM> handleTimer(MessageCollector collector, TaskCoordinator coordinator) {
+    return Collections.emptyList();
   }
 
   /**
-   * Helper method to propagate the output of this operator to all registered downstream operators.
-   *
-   * This method <b>must</b> be called from {@link #onNext} and {@link #onTimer}
-   * to propagate the operator output correctly.
+   * Get the {@link OperatorSpec} for this {@link OperatorImpl}.
    *
-   * @param outputMessage  output message
-   * @param collector  the {@link MessageCollector} in the context
-   * @param coordinator  the {@link TaskCoordinator} in the context
+   * @return the {@link OperatorSpec} for this {@link OperatorImpl}
    */
-  void propagateResult(RM outputMessage, MessageCollector collector, TaskCoordinator coordinator) {
-    nextOperators.forEach(sub -> sub.onNext(outputMessage, collector, coordinator));
+  protected abstract OperatorSpec<RM> getOperatorSpec();
+
+  private HighResolutionClock createHighResClock(Config config) {
+    if (new MetricsConfig(config).getMetricsTimerEnabled()) {
+      return System::nanoTime;
+    } else {
+      return () -> 0;
+    }
   }
 }

samza-core/src/main/java/org/apache/samza/operators/impl/OperatorImplGraph.java

@@ -104,20 +104,22 @@ public Collection<RootOperatorImpl> getAllRootOperators() {
    * creates the corresponding DAG of {@link OperatorImpl}s, and returns its root {@link RootOperatorImpl} node.
    *
    * @param source  the input {@link MessageStreamImpl} to instantiate {@link OperatorImpl}s for
-   * @param <M>  the type of messagess in the {@code source} {@link MessageStreamImpl}
    * @param config  the {@link Config} required to instantiate operators
    * @param context  the {@link TaskContext} required to instantiate operators
+   * @param <M>  the type of messages in the {@code source} {@link MessageStreamImpl}
    * @return  root node for the {@link OperatorImpl} DAG
    */
-  private <M> RootOperatorImpl<M> createOperatorImpls(MessageStreamImpl<M> source, Config config, TaskContext context) {
+  private <M> RootOperatorImpl<M> createOperatorImpls(MessageStreamImpl<M> source,
+      Config config, TaskContext context) {
     // since the source message stream might have multiple operator specs registered on it,
     // create a new root node as a single point of entry for the DAG.
     RootOperatorImpl<M> rootOperator = new RootOperatorImpl<>();
+    rootOperator.init(config, context);
     // create the pipeline/topology starting from the source
     source.getRegisteredOperatorSpecs().forEach(registeredOperator -> {
-        // pass in the source and context s.t. stateful stream operators can initialize their stores
+        // pass in the context so that operator implementations can initialize their functions
         OperatorImpl<M, ?> operatorImpl =
-            createAndRegisterOperatorImpl(registeredOperator, source, config, context);
+            createAndRegisterOperatorImpl(registeredOperator, config, context);
         rootOperator.registerNextOperator(operatorImpl);
       });
     return rootOperator;
@@ -127,27 +129,26 @@ private <M> RootOperatorImpl<M> createOperatorImpls(MessageStreamImpl<M> source,
    * Helper method to recursively traverse the {@link OperatorSpec} DAG and instantiate and link the corresponding
    * {@link OperatorImpl}s.
    *
-   * @param operatorSpec  the operatorSpec registered with the {@code source}
-   * @param source  the source {@link MessageStreamImpl}
-   * @param <M>  type of input message
+   * @param operatorSpec  the operatorSpec to create the {@link OperatorImpl} for
    * @param config  the {@link Config} required to instantiate operators
    * @param context  the {@link TaskContext} required to instantiate operators
+   * @param <M>  type of input message
    * @return  the operator implementation for the operatorSpec
    */
   private <M> OperatorImpl<M, ?> createAndRegisterOperatorImpl(OperatorSpec operatorSpec,
-      MessageStreamImpl<M> source, Config config, TaskContext context) {
+      Config config, TaskContext context) {
     if (!operatorImpls.containsKey(operatorSpec)) {
-      OperatorImpl<M, ?> operatorImpl = createOperatorImpl(source, operatorSpec, config, context);
+      OperatorImpl<M, ?> operatorImpl = createOperatorImpl(operatorSpec, config, context);
       if (operatorImpls.putIfAbsent(operatorSpec, operatorImpl) == null) {
         // this is the first time we've added the operatorImpl corresponding to the operatorSpec,
         // so traverse and initialize and register the rest of the DAG.
         // initialize the corresponding operator function
-        operatorSpec.init(config, context);
+        operatorImpl.init(config, context);
         MessageStreamImpl nextStream = operatorSpec.getNextStream();
         if (nextStream != null) {
           Collection<OperatorSpec> registeredSpecs = nextStream.getRegisteredOperatorSpecs();
           registeredSpecs.forEach(registeredSpec -> {
-              OperatorImpl subImpl = createAndRegisterOperatorImpl(registeredSpec, nextStream, config, context);
+              OperatorImpl subImpl = createAndRegisterOperatorImpl(registeredSpec, config, context);
               operatorImpl.registerNextOperator(subImpl);
             });
         }
@@ -163,24 +164,21 @@ private <M> RootOperatorImpl<M> createOperatorImpls(MessageStreamImpl<M> source,
   /**
    * Creates a new {@link OperatorImpl} instance for the provided {@link OperatorSpec}.
    *
-   * @param source  the source {@link MessageStreamImpl}
-   * @param <M>  type of input message
    * @param operatorSpec  the immutable {@link OperatorSpec} definition.
    * @param config  the {@link Config} required to instantiate operators
    * @param context  the {@link TaskContext} required to instantiate operators
+   * @param <M>  type of input message
    * @return  the {@link OperatorImpl} implementation instance
    */
-  private <M> OperatorImpl<M, ?> createOperatorImpl(MessageStreamImpl<M> source,
-      OperatorSpec operatorSpec, Config config, TaskContext context) {
+  private <M> OperatorImpl<M, ?> createOperatorImpl(OperatorSpec operatorSpec, Config config, TaskContext context) {
     if (operatorSpec instanceof StreamOperatorSpec) {
-      StreamOperatorSpec<M, ?> streamOpSpec = (StreamOperatorSpec<M, ?>) operatorSpec;
-      return new StreamOperatorImpl<>(streamOpSpec, source, config, context);
+      return new StreamOperatorImpl<>((StreamOperatorSpec<M, ?>) operatorSpec, config, context);
     } else if (operatorSpec instanceof SinkOperatorSpec) {
       return new SinkOperatorImpl<>((SinkOperatorSpec<M>) operatorSpec, config, context);
     } else if (operatorSpec instanceof WindowOperatorSpec) {
       return new WindowOperatorImpl((WindowOperatorSpec<M, ?, ?>) operatorSpec, clock);
     } else if (operatorSpec instanceof PartialJoinOperatorSpec) {
-      return new PartialJoinOperatorImpl<>((PartialJoinOperatorSpec) operatorSpec, source, config, context, clock);
+      return new PartialJoinOperatorImpl<>((PartialJoinOperatorSpec) operatorSpec, config, context, clock);
     }
     throw new IllegalArgumentException(
         String.format("Unsupported OperatorSpec: %s", operatorSpec.getClass().getName()));

samza-core/src/main/java/org/apache/samza/operators/impl/PartialJoinOperatorImpl.java

@@ -18,12 +18,11 @@
  */
 package org.apache.samza.operators.impl;
 
-import java.util.ArrayList;
-import java.util.List;
 import org.apache.samza.config.Config;
-import org.apache.samza.operators.MessageStreamImpl;
+import org.apache.samza.metrics.Counter;
 import org.apache.samza.operators.functions.PartialJoinFunction;
 import org.apache.samza.operators.functions.PartialJoinFunction.PartialJoinMessage;
+import org.apache.samza.operators.spec.OperatorSpec;
 import org.apache.samza.operators.spec.PartialJoinOperatorSpec;
 import org.apache.samza.storage.kv.Entry;
 import org.apache.samza.storage.kv.KeyValueIterator;
@@ -35,6 +34,11 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.List;
+
 /**
  * Implementation of a {@link PartialJoinOperatorSpec} that joins messages of type {@code M} in this stream
  * with buffered messages of type {@code JM} in the other stream.
@@ -47,35 +51,45 @@ class PartialJoinOperatorImpl<K, M, JM, RM> extends OperatorImpl<M, RM> {
 
   private static final Logger LOGGER = LoggerFactory.getLogger(PartialJoinOperatorImpl.class);
 
+  private final PartialJoinOperatorSpec<K, M, JM, RM> partialJoinOpSpec;
   private final PartialJoinFunction<K, M, JM, RM> thisPartialJoinFn;
   private final PartialJoinFunction<K, JM, M, RM> otherPartialJoinFn;
   private final long ttlMs;
-  private final int opId;
   private final Clock clock;
 
-  PartialJoinOperatorImpl(PartialJoinOperatorSpec<K, M, JM, RM> partialJoinOperatorSpec, MessageStreamImpl<M> source,
+  private Counter keysRemoved;
+
+  PartialJoinOperatorImpl(PartialJoinOperatorSpec<K, M, JM, RM> partialJoinOpSpec,
       Config config, TaskContext context, Clock clock) {
-    this.thisPartialJoinFn = partialJoinOperatorSpec.getThisPartialJoinFn();
-    this.otherPartialJoinFn = partialJoinOperatorSpec.getOtherPartialJoinFn();
-    this.ttlMs = partialJoinOperatorSpec.getTtlMs();
-    this.opId = partialJoinOperatorSpec.getOpId();
+    this.partialJoinOpSpec = partialJoinOpSpec;
+    this.thisPartialJoinFn = partialJoinOpSpec.getThisPartialJoinFn();
+    this.otherPartialJoinFn = partialJoinOpSpec.getOtherPartialJoinFn();
+    this.ttlMs = partialJoinOpSpec.getTtlMs();
     this.clock = clock;
   }
 
   @Override
-  public void onNext(M message, MessageCollector collector, TaskCoordinator coordinator) {
+  protected void handleInit(Config config, TaskContext context) {
+    keysRemoved = context.getMetricsRegistry()
+        .newCounter(OperatorImpl.class.getName(), this.partialJoinOpSpec.getOpName() + "-keys-removed");
+    this.thisPartialJoinFn.init(config, context);
+  }
+
+  @Override
+  public Collection<RM> handleMessage(M message, MessageCollector collector, TaskCoordinator coordinator) {
     K key = thisPartialJoinFn.getKey(message);
     thisPartialJoinFn.getState().put(key, new PartialJoinMessage<>(message, clock.currentTimeMillis()));
     PartialJoinMessage<JM> otherMessage = otherPartialJoinFn.getState().get(key);
     long now = clock.currentTimeMillis();
     if (otherMessage != null && otherMessage.getReceivedTimeMs() > now - ttlMs) {
       RM joinResult = thisPartialJoinFn.apply(message, otherMessage.getMessage());
-      this.propagateResult(joinResult, collector, coordinator);
+      return Collections.singletonList(joinResult);
     }
+    return Collections.emptyList();
   }
 
   @Override
-  public void onTimer(MessageCollector collector, TaskCoordinator coordinator) {
+  public Collection<RM> handleTimer(MessageCollector collector, TaskCoordinator coordinator) {
     long now = clock.currentTimeMillis();
 
     KeyValueStore<K, PartialJoinMessage<M>> thisState = thisPartialJoinFn.getState();
@@ -87,14 +101,18 @@ public void onTimer(MessageCollector collector, TaskCoordinator coordinator) {
       if (entry.getValue().getReceivedTimeMs() < now - ttlMs) {
         keysToRemove.add(entry.getKey());
       } else {
-        break;
+        break; // InternalInMemoryStore uses a LinkedHashMap and will return entries in insertion order
       }
     }
 
     iterator.close();
     thisState.deleteAll(keysToRemove);
-
-    LOGGER.info("Operator ID {} onTimer self time: {} ms", opId, clock.currentTimeMillis() - now);
+    keysRemoved.inc(keysToRemove.size());
+    return Collections.emptyList();
   }
 
+  @Override
+  protected OperatorSpec<RM> getOperatorSpec() {
+    return partialJoinOpSpec;
+  }
 }