Concurrency limiter from Netflix to avoid throttling storage

pom.xml

@@ -64,6 +64,7 @@
     <!-- Be careful to set this as a provided dep, so that it doesn't interfere with dependencies
          from other projects. For example, cassandra and spring boot set guava versions -->
     <guava.version>19.0</guava.version>
+    <netflix.concurrency.limits.version>0.0.49</netflix.concurrency.limits.version>
     <junit.version>4.12</junit.version>
     <powermock.version>2.0.0-beta.5</powermock.version>
     <mockito.version>2.21.0</mockito.version>
@@ -396,6 +397,14 @@
         <artifactId>brave-instrumentation-cassandra-driver</artifactId>
         <version>0.8.0</version>
       </dependency>
+
+      <!-- Collector concurrency limits -->
+      <dependency>
+        <groupId>com.netflix.concurrency-limits</groupId>
+        <artifactId>concurrency-limits-core</artifactId>
+        <version>${netflix.concurrency.limits.version}</version>
+      </dependency>
+
     </dependencies>
   </dependencyManagement>
 

zipkin-autoconfigure/collector-kafka/src/main/java/zipkin2/autoconfigure/collector/kafka/ZipkinKafkaCollectorAutoConfiguration.java

@@ -13,6 +13,7 @@
  */
 package zipkin2.autoconfigure.collector.kafka;
 
+import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.boot.context.properties.EnableConfigurationProperties;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Condition;
@@ -22,6 +23,7 @@
 import org.springframework.core.type.AnnotatedTypeMetadata;
 import zipkin2.collector.CollectorMetrics;
 import zipkin2.collector.CollectorSampler;
+import zipkin2.collector.ConcurrencyLimiter;
 import zipkin2.collector.kafka.KafkaCollector;
 import zipkin2.storage.StorageComponent;
 
@@ -36,11 +38,12 @@ class ZipkinKafkaCollectorAutoConfiguration { // makes simple type name unique f
 
   @Bean(initMethod = "start")
   KafkaCollector kafka(
-      ZipkinKafkaCollectorProperties properties,
-      CollectorSampler sampler,
-      CollectorMetrics metrics,
-      StorageComponent storage) {
-    return properties.toBuilder().sampler(sampler).metrics(metrics).storage(storage).build();
+    ZipkinKafkaCollectorProperties properties,
+    CollectorSampler sampler,
+    CollectorMetrics metrics,
+    StorageComponent storage,
+    @Autowired(required = false) ConcurrencyLimiter limiter) {
+    return properties.toBuilder().sampler(sampler).metrics(metrics).storage(storage).limiter(limiter).build();
   }
 
   /**

zipkin-autoconfigure/collector-kafka/src/main/java/zipkin2/autoconfigure/collector/kafka/ZipkinKafkaCollectorProperties.java

@@ -28,6 +28,7 @@ class ZipkinKafkaCollectorProperties {
   private String topic;
   /** Number of Kafka consumer threads to run. */
   private Integer streams;
+
   /** Additional Kafka consumer configuration. */
   private Map<String, String> overrides = new LinkedHashMap<>();
 
@@ -78,10 +79,12 @@ public KafkaCollector.Builder toBuilder() {
     if (topic != null) result.topic(topic);
     if (streams != null) result.streams(streams);
     if (overrides != null) result.overrides(overrides);
+
     return result;
   }
 
   private static String emptyToNull(String s) {
     return "".equals(s) ? null : s;
   }
+
 }

zipkin-autoconfigure/collector-kafka08/src/main/java/zipkin2/autoconfigure/collector/kafka08/ZipkinKafka08CollectorAutoConfiguration.java

@@ -13,12 +13,14 @@
  */
 package zipkin2.autoconfigure.collector.kafka08;
 
+import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.boot.context.properties.EnableConfigurationProperties;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Conditional;
 import org.springframework.context.annotation.Configuration;
 import zipkin2.collector.CollectorMetrics;
 import zipkin2.collector.CollectorSampler;
+import zipkin2.collector.ConcurrencyLimiter;
 import zipkin2.collector.kafka08.KafkaCollector;
 import zipkin2.storage.StorageComponent;
 
@@ -40,9 +42,10 @@ KafkaCollector kafka(
       ZipkinKafkaCollectorProperties kafka,
       CollectorSampler sampler,
       CollectorMetrics metrics,
-      StorageComponent storage) {
+      StorageComponent storage,
+      @Autowired(required = false) ConcurrencyLimiter limiter) {
     final KafkaCollector result =
-        kafka.toBuilder().sampler(sampler).metrics(metrics).storage(storage).build();
+        kafka.toBuilder().sampler(sampler).metrics(metrics).storage(storage).limiter(limiter).build();
 
     // don't use @Bean(initMethod = "start") as it can crash the process if zookeeper is down
     Thread start =

zipkin-autoconfigure/collector-kafka08/src/main/java/zipkin2/autoconfigure/collector/kafka08/ZipkinKafkaCollectorProperties.java

@@ -77,11 +77,11 @@ public void setOverrides(Map<String, String> overrides) {
 
   public KafkaCollector.Builder toBuilder() {
     return KafkaCollector.builder()
-        .topic(topic)
-        .zookeeper(zookeeper)
-        .groupId(groupId)
-        .streams(streams)
-        .maxMessageSize(maxMessageSize)
-        .overrides(overrides);
+      .topic(topic)
+      .zookeeper(zookeeper)
+      .groupId(groupId)
+      .streams(streams)
+      .maxMessageSize(maxMessageSize)
+      .overrides(overrides);
   }
 }

zipkin-autoconfigure/collector-rabbitmq/src/main/java/zipkin2/autoconfigure/collector/rabbitmq/ZipkinRabbitMQCollectorAutoConfiguration.java

@@ -16,6 +16,8 @@
 import java.net.URISyntaxException;
 import java.security.KeyManagementException;
 import java.security.NoSuchAlgorithmException;
+
+import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.boot.context.properties.EnableConfigurationProperties;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Condition;
@@ -25,6 +27,7 @@
 import org.springframework.core.type.AnnotatedTypeMetadata;
 import zipkin2.collector.CollectorMetrics;
 import zipkin2.collector.CollectorSampler;
+import zipkin2.collector.ConcurrencyLimiter;
 import zipkin2.collector.rabbitmq.RabbitMQCollector;
 import zipkin2.storage.StorageComponent;
 
@@ -36,12 +39,13 @@ class ZipkinRabbitMQCollectorAutoConfiguration {
 
   @Bean(initMethod = "start")
   RabbitMQCollector rabbitMq(
-      ZipkinRabbitMQCollectorProperties properties,
-      CollectorSampler sampler,
-      CollectorMetrics metrics,
-      StorageComponent storage)
+    ZipkinRabbitMQCollectorProperties properties,
+    CollectorSampler sampler,
+    CollectorMetrics metrics,
+    StorageComponent storage,
+    @Autowired(required = false) ConcurrencyLimiter limiter)
       throws NoSuchAlgorithmException, KeyManagementException, URISyntaxException {
-    return properties.toBuilder().sampler(sampler).metrics(metrics).storage(storage).build();
+    return properties.toBuilder().sampler(sampler).metrics(metrics).storage(storage).limiter(limiter).build();
   }
 
   /**

zipkin-collector/core/pom.xml

@@ -32,6 +32,13 @@
   </properties>
 
   <dependencies>
+
+    <!-- Concurrency limits -->
+    <dependency>
+      <groupId>com.netflix.concurrency-limits</groupId>
+      <artifactId>concurrency-limits-core</artifactId>
+    </dependency>
+
     <dependency>
       <groupId>${project.groupId}</groupId>
       <artifactId>zipkin</artifactId>

zipkin-collector/core/src/main/java/zipkin2/collector/Collector.java

@@ -13,6 +13,7 @@
  */
 package zipkin2.collector;
 
+import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.Iterator;
@@ -48,6 +49,7 @@ public static final class Builder {
     StorageComponent storage = null;
     CollectorSampler sampler = null;
     CollectorMetrics metrics = null;
+    ConcurrencyLimiter limiter = null;
 
     Builder(Logger logger) {
       this.logger = logger;
@@ -74,6 +76,13 @@ public Builder sampler(CollectorSampler sampler) {
       return this;
     }
 
+    /** @see {@link CollectorComponent.Builder#limiter(ConcurrencyLimiter)} */
+    public Builder limiter(ConcurrencyLimiter limiter) {
+      if (limiter == null) throw new NullPointerException("limiter == null");
+      this.limiter = limiter;
+      return this;
+    }
+
     public Collector build() {
       return new Collector(this);
     }
@@ -83,6 +92,7 @@ public Collector build() {
   final CollectorMetrics metrics;
   final CollectorSampler sampler;
   final StorageComponent storage;
+  final ConcurrencyLimiter limiter;
 
   Collector(Builder builder) {
     if (builder.logger == null) throw new NullPointerException("logger == null");
@@ -91,6 +101,7 @@ public Collector build() {
     if (builder.storage == null) throw new NullPointerException("storage == null");
     this.storage = builder.storage;
     this.sampler = builder.sampler == null ? CollectorSampler.ALWAYS_SAMPLE : builder.sampler;
+    this.limiter = builder.limiter;
   }
 
   public void accept(List<Span> spans, Callback<Void> callback) {
@@ -147,7 +158,20 @@ List<Span> decodeList(BytesDecoder<Span> decoder, byte[] serialized) {
   }
 
   void record(List<Span> sampled, Callback<Void> callback) {
-    storage.spanConsumer().accept(sampled).enqueue(callback);
+    if(limiter != null) {
+
+      limiter.execute(() -> {
+        try {
+          storage.spanConsumer().accept(sampled).execute();
+          callback.onSuccess(null);
+        } catch (IOException e) {
+          callback.onError(e);
+        }
+      });
+
+    } else {
+      storage.spanConsumer().accept(sampled).enqueue(callback);
+    }
   }
 
   String idString(Span span) {

zipkin-collector/core/src/main/java/zipkin2/collector/CollectorComponent.java

@@ -52,6 +52,15 @@ public abstract static class Builder {
      */
     public abstract Builder sampler(CollectorSampler sampler);
 
+    /**
+     * Apply limits based on concurrency during sample record. Storage signals over capacity with a
+     * RejectedExecutionException.
+     *
+     * @param limiter limiter
+     * @return fluent builder instance.
+     */
+    public abstract Builder limiter(ConcurrencyLimiter limiter);
+
     public abstract CollectorComponent build();
   }
 }

zipkin-collector/core/src/main/java/zipkin2/collector/ConcurrencyLimiter.java

@@ -0,0 +1,57 @@
+/*
+ * Copyright 2015-2018 The OpenZipkin Authors
+ *
+ * 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 zipkin2.collector;
+
+import com.netflix.concurrency.limits.Limiter;
+import com.netflix.concurrency.limits.executors.BlockingAdaptiveExecutor;
+
+import java.util.concurrent.Executor;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+
+public class ConcurrencyLimiter {
+
+  private int threads = 1;
+  private ExecutorService pool;
+  private Executor executor;
+
+  public void setThreads(int threads) {
+    this.threads = threads;
+  }
+
+  public void setLimiter(Limiter<Void> limiter) {
+    if(threads > 1) {
+      pool = Executors.newFixedThreadPool(threads);
+    } else {
+      pool = Executors.newSingleThreadExecutor();
+    }
+    this.executor = new BlockingAdaptiveExecutor(limiter, pool);
+  }
+
+  public void execute(Runnable call) {
+    executor.execute(call);
+  }
+
+  public void close() {
+    if (pool == null) return;
+    pool.shutdownNow();
+    try {
+      pool.awaitTermination(1, TimeUnit.SECONDS);
+    } catch (InterruptedException e) {
+      // do nothing
+    }
+  }
+
+}

zipkin-collector/kafka/src/main/java/zipkin2/collector/kafka/KafkaCollector.java

@@ -27,10 +27,7 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import zipkin2.CheckResult;
-import zipkin2.collector.Collector;
-import zipkin2.collector.CollectorComponent;
-import zipkin2.collector.CollectorMetrics;
-import zipkin2.collector.CollectorSampler;
+import zipkin2.collector.*;
 import zipkin2.storage.SpanConsumer;
 import zipkin2.storage.StorageComponent;
 
@@ -60,6 +57,7 @@ public static final class Builder extends CollectorComponent.Builder {
     CollectorMetrics metrics = CollectorMetrics.NOOP_METRICS;
     String topic = "zipkin";
     int streams = 1;
+    ConcurrencyLimiter limiter;
 
     @Override
     public Builder storage(StorageComponent storage) {
@@ -81,6 +79,12 @@ public Builder metrics(CollectorMetrics metrics) {
       return this;
     }
 
+    @Override
+    public Builder limiter(ConcurrencyLimiter limiter) {
+      this.limiter = limiter;
+      return this;
+    }
+
     /**
      * Topic zipkin spans will be consumed from. Defaults to "zipkin". Multiple topics may be
      * specified if comma delimited.
@@ -148,9 +152,11 @@ public KafkaCollector build() {
   }
 
   final LazyKafkaWorkers kafkaWorkers;
+  final ConcurrencyLimiter limiter;
 
   KafkaCollector(Builder builder) {
     kafkaWorkers = new LazyKafkaWorkers(builder);
+    limiter = builder.limiter;
   }
 
   @Override
@@ -173,6 +179,9 @@ public CheckResult check() {
   @Override
   public void close() {
     kafkaWorkers.close();
+    if(limiter != null) {
+      limiter.close();
+    }
   }
 
   static final class LazyKafkaWorkers {