content: /*
   * 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.beam.sdk.extensions.euphoria.core.testkit.accumulators;
  
  import java.io.ObjectStreamException;
  import java.time.Duration;
  import java.util.HashMap;
  import java.util.Map;
  import java.util.concurrent.ConcurrentHashMap;
  import java.util.concurrent.ConcurrentMap;
  import org.apache.beam.sdk.extensions.euphoria.core.client.accumulators.Accumulator;
  import org.apache.beam.sdk.extensions.euphoria.core.client.accumulators.AccumulatorProvider;
  import org.apache.beam.sdk.extensions.euphoria.core.client.accumulators.Counter;
  import org.apache.beam.sdk.extensions.euphoria.core.client.accumulators.Histogram;
  import org.apache.beam.sdk.extensions.euphoria.core.client.accumulators.Timer;
  
  /**
   * An accumulator provider gathering accumulators in-memory.
   *
   * <p>Safe for use for unit testing purposes.
   */
  public class SingleJvmAccumulatorProvider implements AccumulatorProvider {
  
    private static final SingleJvmAccumulatorProvider INSTANCE = new SingleJvmAccumulatorProvider();
    private final ConcurrentMap<String, Accumulator> accs = new ConcurrentHashMap<>();
  
    private SingleJvmAccumulatorProvider() {}
  
    @SuppressWarnings("unchecked")
    private static <T> T assertType(String name, Class<T> expectedType, Accumulator actualAcc) {
      if (actualAcc.getClass() != expectedType) {
        // ~ provide a nice message (that's why we don't simply use `expectedType.cast(..)`)
        throw new IllegalStateException(
            "Ambiguously named accumulators! Got "
                + actualAcc.getClass()
                + " for "
                + name
                + " but expected "
                + expectedType
                + "!");
      }
      return (T) actualAcc;
    }
  
    @Override
    public Counter getCounter(String name) {
      return assertType(name, LongCounter.class, accs.computeIfAbsent(name, s -> new LongCounter()));
    }
  
    @Override
    public Counter getCounter(String namespace, String name) {
      return getCounter(name);
    }
  
    @Override
    public Histogram getHistogram(String name) {
      return assertType(
          name, LongHistogram.class, accs.computeIfAbsent(name, s -> new LongHistogram()));
    }
  
    @Override
    public Histogram getHistogram(String namespace, String name) {
      return getHistogram(name);
    }
  
    @Override
    public Timer getTimer(String name) {
      return assertType(
          name, NanosecondTimer.class, accs.computeIfAbsent(name, s -> new NanosecondTimer()));
    }
  
    void clear() {
      accs.clear();
    }
  
    <V, T extends Snapshotable<V>> Map<String, V> getSnapshots(Class<T> type) {
      HashMap<String, V> m = new HashMap<>();
      accs.forEach(
          (name, accumulator) -> {
            if (type.isAssignableFrom(accumulator.getClass())) {
              @SuppressWarnings("unchecked")
              T acc = (T) accumulator;
              m.put(name, acc.getSnapshot());
            }
          });
      return m;
    }
  
    // ~ -----------------------------------------------------------------------
  
    /** Accumulator provider factory. */
    public static final class Factory implements AccumulatorProvider.Factory, SnapshotProvider {
  
      private static final Factory INSTANCE = new Factory();
  
      private Factory() {}
  
      /**
       * Before running another test, clear method must be called to delete metrics snapshots.
       *
       * @return SingleJvmAccumulatorFactory
       */
      public static Factory get() {
        return INSTANCE;
      }
  
      public void clear() {
        providerInstance().clear();
      }
  
      @Override
      public Map<String, Long> getCounterSnapshots() {
        return providerInstance().getSnapshots(LongCounter.class);
      }
  
      @Override
      public Map<String, Map<Long, Long>> getHistogramSnapshots() {
        return providerInstance().getSnapshots(LongHistogram.class);
      }
  
      @Override
      public Map<String, Map<Duration, Long>> getTimerSnapshots() {
        return providerInstance().getSnapshots(NanosecondTimer.class);
      }
  
      @Override
      public AccumulatorProvider create() {
        return providerInstance();
      }
  
      private SingleJvmAccumulatorProvider providerInstance() {
        return SingleJvmAccumulatorProvider.INSTANCE;
      }
  
      private Object readResolve() throws ObjectStreamException {
        return INSTANCE;
      }
    }
  }
  