From 645845e88c9947bdeb006da6c450291ee243fcff Mon Sep 17 00:00:00 2001
From: Marcin Grzejszczak <mgrzejszczak@pivotal.io>
Date: Fri, 29 Apr 2016 16:00:59 +0200
Subject: [PATCH] Upgraded the percentage based sampler algorithm

---
 .../sampler/PercentageBasedSampler.java       | 53 ++++++++++++++-----
 1 file changed, 39 insertions(+), 14 deletions(-)

diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/sampler/PercentageBasedSampler.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/sampler/PercentageBasedSampler.java
index 5c879cd1c..2c931d953 100644
--- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/sampler/PercentageBasedSampler.java
+++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/sampler/PercentageBasedSampler.java
@@ -1,5 +1,9 @@
 package org.springframework.cloud.sleuth.sampler;
 
+import java.util.BitSet;
+import java.util.Random;
+import java.util.concurrent.atomic.AtomicInteger;
+
 import org.springframework.cloud.sleuth.Sampler;
 import org.springframework.cloud.sleuth.Span;
 
@@ -21,33 +25,54 @@ import org.springframework.cloud.sleuth.Span;
  */
 public class PercentageBasedSampler implements Sampler {
 
-	private final int outOf100;
-
-	private int i = 0; // guarded by this
-	private boolean skipping = false; // guarded by this
+	private final AtomicInteger counter = new AtomicInteger(0);
+	private final BitSet sampleDecisions;
+	private final SamplerProperties configuration;
 
 	public PercentageBasedSampler(SamplerProperties configuration) {
-		this.outOf100 = (int) (configuration.getPercentage() * 100.0f);;
+		int outOf100 = (int) (configuration.getPercentage() * 100.0f);;
+		this.sampleDecisions = randomBitSet(100, outOf100, new Random());
+		this.configuration = configuration;
 	}
 
 	@Override
 	public boolean isSampled(Span currentSpan) {
-		if (this.outOf100 == 0 || currentSpan == null) {
+		if (this.configuration.getPercentage() == 0 || currentSpan == null) {
 			return false;
-		} else if (this.outOf100 == 100) {
+		} else if (this.configuration.getPercentage() == 100) {
 			return true;
 		}
 		synchronized (this) {
-			boolean result = !this.skipping;
-			this.i = this.i + 1;
-			if (this.i == this.outOf100) {
-				this.skipping = true;
-			} else if (this.i == 100) {
-				this.i = 0;
-				this.skipping = false;
+			final int i = this.counter.getAndIncrement();
+			boolean result = this.sampleDecisions.get(i);
+			if (i == 99) {
+				this.counter.set(0);
 			}
 			return result;
 		}
 	}
 
+	/**
+	 * Reservoir sampling algorithm borrowed from Stack Overflow.
+	 *
+	 * http://stackoverflow.com/questions/12817946/generate-a-random-bitset-with-n-1s
+	 */
+	static BitSet randomBitSet(int size, int cardinality, Random rnd) {
+		BitSet result = new BitSet(size);
+		int[] chosen = new int[cardinality];
+		int i;
+		for (i = 0; i < cardinality; ++i) {
+			chosen[i] = i;
+			result.set(i);
+		}
+		for (; i < size; ++i) {
+			int j = rnd.nextInt(i + 1);
+			if (j < cardinality) {
+				result.clear(chosen[j]);
+				result.set(i);
+				chosen[j] = i;
+			}
+		}
+		return result;
+	}
 }