/*
 * Copyright (C) 2011 Google Inc.
 *
 * Licensed 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 com.google.caliper.worker;

import static com.google.common.base.Preconditions.checkArgument;

import com.google.caliper.Benchmark;
import com.google.caliper.model.Measurement;
import com.google.caliper.model.Value;
import com.google.common.base.Objects;
import com.google.common.collect.ImmutableList;
import com.google.inject.Inject;
import com.google.monitoring.runtime.instrumentation.AllocationRecorder;
import com.google.monitoring.runtime.instrumentation.Sampler;

import java.lang.reflect.Method;
import java.util.Arrays;
import java.util.Map;
import java.util.Random;

/**
 * The {@link Worker} for the {@code AllocationInstrument}.  This class invokes the benchmark method
 * a few times, with varying numbers of reps, and computes the number of object allocations and the
 * total size of those allocations.
 */
public final class AllocationWorker implements Worker {
  private static final int MAX_BASELINE_REPS = 5;
  private static final int MAX_REPS_ABOVE_BASELINE = 100;

  private final Random random;

  private int allocationCount;
  private long allocationSize;
  private boolean recordAllocations = false;

  @Inject AllocationWorker(Random random) {
    this.random = random;
    AllocationRecorder.addSampler(
        new Sampler() {
          @Override
          public void sampleAllocation(int arrayCount, String desc, Object newObj, long size) {
            if (recordAllocations) {
              allocationCount++;
              allocationSize += size;
            }
          }
        });
  }

  @Override public synchronized void measure(Benchmark benchmark, String methodName,
      Map<String, String> options, WorkerEventLog log) throws Exception {
    // do one initial measurement and throw away its results
    log.notifyWarmupPhaseStarting();
    measureAllocations(benchmark, methodName, 1);

    log.notifyMeasurementPhaseStarting();
    while (true) {
      log.notifyMeasurementStarting();
      // [1, 5]
      int baselineReps = random.nextInt(MAX_BASELINE_REPS) + 1;
      AllocationStats baseline = measureAllocations(benchmark, methodName, baselineReps);
      // (baseline, baseline + MAX_REPS_ABOVE_BASELINE]
      int measurementReps = baselineReps + random.nextInt(MAX_REPS_ABOVE_BASELINE) + 1;
      AllocationStats measurement = measureAllocations(benchmark, methodName, measurementReps);
      try {
        AllocationStats diff = diffMeasurements(baseline, measurement);
        log.notifyMeasurementEnding(ImmutableList.of(
            new Measurement.Builder()
                .value(Value.create(diff.allocationCount, ""))
                .description("objects")
                .weight(diff.reps)
                .build(),
            new Measurement.Builder()
                .value(Value.create(diff.allocationSize, "B"))
                .weight(diff.reps)
                .description("bytes")
                .build()));
      } catch (IllegalStateException e) {
        // log the failure, but just keep trying to measure
        log.notifyMeasurementFailure(e);
      }
    }
  }

  /**
   * Computes and returns the difference between these two measurements. The {@code initial}
   * measurement must have a lower weight (fewer reps) than the {@code second} measurement.
   */
  private AllocationStats diffMeasurements(AllocationStats baseline, AllocationStats measurement) {
    try {
      return new AllocationStats(measurement.allocationCount - baseline.allocationCount,
          measurement.allocationSize - baseline.allocationSize,
          measurement.reps - baseline.reps);
    } catch (IllegalArgumentException e) {
      throw new IllegalStateException(String.format(
          "The difference between the baseline (%s) and the measurement (%s) is invalid",
              baseline, measurement), e);
    }
  }

  static Method findMethod(Class<? extends Benchmark> benchmarkClass, String name) {
    for (Method method : benchmarkClass.getDeclaredMethods()) {
      Class<?>[] parameterTypes = method.getParameterTypes();
      if (method.getName().equals(name)
          && (Arrays.equals(parameterTypes, new Class<?>[] {int.class})
                || Arrays.equals(parameterTypes, new Class<?>[] {long.class}))) {
        return method;
      }
    }
    throw new IllegalArgumentException(
        String.format("no method named %s on %s that takes either an int or a long",
            name, benchmarkClass));
  }

  private synchronized AllocationStats measureAllocations(
      Benchmark benchmark, String methodName, int reps) throws Exception {
    Method method = findMethod(benchmark.getClass(), "time" + methodName);
    clearAccumulatedStats();
    // do the Integer boxing and the creation of the Object[] outside of the record block, so that
    // our internal allocations aren't counted in the benchmark's allocations.
    Object[] args = {reps};
    recordAllocations = true;
    method.invoke(benchmark, args);
    recordAllocations = false;

    return new AllocationStats(allocationCount, allocationSize, reps);
  }

  private synchronized void clearAccumulatedStats() {
    recordAllocations = false;
    allocationCount = 0;
    allocationSize = 0;
  }

  static class AllocationStats {
    final int allocationCount;
    final long allocationSize;
    final int reps;

    AllocationStats(int allocationCount, long allocationSize, int reps) {
      checkArgument(allocationCount >= 0, "allocationCount (%s) was negative", allocationCount);
      this.allocationCount = allocationCount;
      checkArgument(allocationSize >= 0, "allocationSize (%s) was negative", allocationSize);
      this.allocationSize = allocationSize;
      checkArgument(reps >= 0, "reps (%s) was negative", reps);
      this.reps = reps;
    }

    @Override public String toString() {
      return Objects.toStringHelper(this)
          .add("allocationCount", allocationCount)
          .add("allocationSize", allocationSize)
          .add("reps", reps)
          .toString();
    }
  }
}

@@ -1173,16 +1173,113 @@
 .Random;
+%0Aimport java.util.concurrent.atomic.AtomicInteger;%0Aimport java.util.concurrent.atomic.AtomicLong;
 %0A%0A/**%0A *
@@ -1728,58 +1728,127 @@
 ate 
+f
 in
-t allocationCount;%0A  private long allocationSize
+al AtomicInteger allocationCount = new AtomicInteger();%0A  private final AtomicLong allocationSize = new AtomicLong()
 ;%0A  
@@ -1855,16 +1855,25 @@
 private 
+volatile 
 boolean 
@@ -2215,10 +2215,26 @@
 ount
-++
+.getAndIncrement()
 ;%0A  
@@ -2263,16 +2263,24 @@
 Size
- += 
+.getAndAdd(
 size
+)
 ;%0A  
@@ -2320,16 +2320,16 @@
 );%0A  %7D%0A%0A
+
   @Overr
@@ -2339,29 +2339,16 @@
  public 
-synchronized 
 void mea
@@ -5191,29 +5191,16 @@
 private 
-synchronized 
 Allocati
@@ -5756,16 +5756,22 @@
 ionCount
+.get()
 , alloca
@@ -5778,16 +5778,22 @@
 tionSize
+.get()
 , reps);
@@ -5812,21 +5812,8 @@
 ate 
-synchronized 
 void
@@ -5889,20 +5889,23 @@
 ionCount
- = 0
+.set(0)
 ;%0A    al
@@ -5920,12 +5920,15 @@
 Size
- = 0
+.set(0)
 ;%0A  

make the AllocationWorker a little more thread-safe.