Add RemoteBaggageRestrictionManager (#229)

jaeger-core/src/main/java/com/uber/jaeger/Span.java

@@ -22,7 +22,6 @@
 
 package com.uber.jaeger;
 
-import com.uber.jaeger.baggage.BaggageSetter;
 import io.opentracing.tag.Tags;
 import java.util.ArrayList;
 import java.util.Collections;
@@ -106,6 +105,12 @@ public String getOperationName() {
     }
   }
 
+  public String getServiceName() {
+    synchronized (this) {
+      return this.getTracer().getServiceName();
+    }
+  }
+
   public List<LogData> getLogs() {
     synchronized (this) {
       if (logs == null) {

jaeger-core/src/main/java/com/uber/jaeger/baggage/BaggageRestrictionManagerProxy.java

@@ -0,0 +1,37 @@
+/*
+ * Copyright (c) 2017, Uber Technologies, Inc
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+
+package com.uber.jaeger.baggage;
+
+import com.uber.jaeger.baggage.http.BaggageRestrictionResponse;
+import com.uber.jaeger.exceptions.BaggageRestrictionManagerException;
+
+import java.util.List;
+
+/**
+ * BaggageRestrictionManagerProxy is an interface for a class that fetches baggage
+ * restrictions for specific service from a remote source i.e. jaeger-agent.
+ */
+public interface BaggageRestrictionManagerProxy {
+  List<BaggageRestrictionResponse> getBaggageRestrictions(String serviceName)
+      throws BaggageRestrictionManagerException;
+}

jaeger-core/src/main/java/com/uber/jaeger/baggage/BaggageSetter.java

@@ -58,7 +58,7 @@ public BaggageSetter(BaggageRestrictionManager restrictionManager, Metrics metri
    * @return       the SpanContext with the baggage set
    */
   public SpanContext setBaggage(Span span, String key, String value) {
-    Restriction restriction = restrictionManager.getRestriction(span.getTracer().getServiceName(), key);
+    Restriction restriction = restrictionManager.getRestriction(span.getServiceName(), key);
     boolean truncated = false;
     String prevItem = null;
     if (!restriction.isKeyAllowed()) {

jaeger-core/src/main/java/com/uber/jaeger/baggage/HttpBaggageRestrictionManagerProxy.java

@@ -0,0 +1,71 @@
+/*
+ * Copyright (c) 2017, Uber Technologies, Inc
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+
+package com.uber.jaeger.baggage;
+
+import static com.uber.jaeger.utils.Http.makeGetRequest;
+
+import com.google.gson.Gson;
+import com.google.gson.JsonSyntaxException;
+import com.google.gson.reflect.TypeToken;
+import com.uber.jaeger.baggage.http.BaggageRestrictionResponse;
+import com.uber.jaeger.exceptions.BaggageRestrictionManagerException;
+
+import java.io.IOException;
+import java.lang.reflect.Type;
+import java.net.URLEncoder;
+import java.util.ArrayList;
+import java.util.List;
+
+public class HttpBaggageRestrictionManagerProxy implements BaggageRestrictionManagerProxy {
+  private static final String DEFAULT_HOST_PORT = "localhost:5778";
+  private final Gson gson = new Gson();
+  private final String hostPort;
+
+  public HttpBaggageRestrictionManagerProxy(String hostPort) {
+    this.hostPort = hostPort != null ? hostPort : DEFAULT_HOST_PORT;
+  }
+
+  List<BaggageRestrictionResponse> parseJson(String json) throws BaggageRestrictionManagerException {
+    try {
+      Type listType = new TypeToken<ArrayList<BaggageRestrictionResponse>>(){}.getType();
+      return gson.fromJson(json, listType);
+    } catch (JsonSyntaxException e) {
+      throw new BaggageRestrictionManagerException("Cannot deserialize json", e);
+    }
+  }
+
+  @Override
+  public List<BaggageRestrictionResponse> getBaggageRestrictions(String serviceName)
+      throws BaggageRestrictionManagerException {
+    String jsonString;
+    try {
+      jsonString =
+          makeGetRequest(
+              "http://" + hostPort + "/baggageRestrictions?service=" + URLEncoder.encode(serviceName, "UTF-8"));
+    } catch (IOException e) {
+      throw new BaggageRestrictionManagerException(
+          "http call to get baggage restriction from local agent failed.", e);
+    }
+    return parseJson(jsonString);
+  }
+}

jaeger-core/src/main/java/com/uber/jaeger/baggage/RemoteBaggageRestrictionManager.java

@@ -0,0 +1,167 @@
+/*
+ * Copyright (c) 2017, Uber Technologies, Inc
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+
+package com.uber.jaeger.baggage;
+
+import com.uber.jaeger.baggage.http.BaggageRestrictionResponse;
+import com.uber.jaeger.exceptions.BaggageRestrictionManagerException;
+import com.uber.jaeger.metrics.Metrics;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Timer;
+import java.util.TimerTask;
+
+/**
+ * NewRestrictionManager returns a BaggageRestrictionManager that polls the agent for the latest
+ * baggage restrictions.
+ */
+public class RemoteBaggageRestrictionManager implements BaggageRestrictionManager {
+  private static final int DEFAULT_REFRESH_INTERVAL_MS = 60000;
+  private static final int DEFAULT_INITIAL_DELAY_MS = 0;
+
+  private final String serviceName;
+  private final BaggageRestrictionManagerProxy proxy;
+  private final Timer pollTimer;
+  private final Metrics metrics;
+  private final boolean denyBaggageOnInitializationFailure;
+  private volatile boolean initialized;
+  private volatile Map<String, Restriction> restrictions = new HashMap<String, Restriction>();
+  private final Restriction invalidRestriction;
+  private final Restriction validRestriction;
+
+  public RemoteBaggageRestrictionManager(
+      String serviceName,
+      BaggageRestrictionManagerProxy proxy,
+      Metrics metrics,
+      boolean denyBaggageOnInitializationFailure
+  ) {
+    this(serviceName, proxy, metrics, denyBaggageOnInitializationFailure, DEFAULT_REFRESH_INTERVAL_MS);
+  }
+
+  public RemoteBaggageRestrictionManager(
+      String serviceName,
+      BaggageRestrictionManagerProxy proxy,
+      Metrics metrics,
+      boolean denyBaggageOnInitializationFailure,
+      int refreshIntervalMs
+  ) {
+    this(serviceName, proxy, metrics, denyBaggageOnInitializationFailure, refreshIntervalMs, DEFAULT_INITIAL_DELAY_MS);
+  }
+
+  /**
+   * Creates a RemoteBaggageRestrictionManager that fetches {@link BaggageRestrictionResponse} from a remote
+   * agent and keeps track of {@link Restriction} for a service.
+   *
+   * {@param initialDelayMs} is only exposed for testing purposes so users can determine when the first call to
+   * remote agent is made. Under normal operations, this RemoteBaggageRestrictionManager will start up and
+   * asynchronously fetch restrictions. If the user wants to know if restrictions are ready, they can check via
+   * isReady().
+   *
+   * @param serviceName restrictions for this service are kept track of.
+   * @param proxy proxy to remote agent.
+   * @param metrics metrics for metrics emission.
+   * @param denyBaggageOnInitializationFailure determines the startup failure mode of RemoteBaggageRestrictionManager.
+   *                                           If DenyBaggageOnInitializationFailure is true,
+   *                                           RemoteBaggageRestrictionManager will not allow any baggage to be written
+   *                                           until baggage restrictions have been retrieved from agent. If
+   *                                           DenyBaggageOnInitializationFailure is false,
+   *                                           RemoteBaggageRestrictionManager will allow any baggage to be written
+   *                                           until baggage restrictions have been retrieved from agent.
+   * @param refreshIntervalMs how often restriction are fetched from remote agent.
+   * @param initialDelayMs delay before first fetch of restrictions.
+   */
+  RemoteBaggageRestrictionManager(
+      String serviceName,
+      BaggageRestrictionManagerProxy proxy,
+      Metrics metrics,
+      boolean denyBaggageOnInitializationFailure,
+      int refreshIntervalMs,
+      int initialDelayMs
+  ) {
+    this.serviceName = serviceName;
+    this.proxy = proxy;
+    this.metrics = metrics;
+    this.denyBaggageOnInitializationFailure = denyBaggageOnInitializationFailure;
+    this.initialized = false;
+    this.invalidRestriction = Restriction.of(false, 0);
+    this.validRestriction = Restriction.of(true, DEFAULT_MAX_VALUE_LENGTH);
+
+    pollTimer = new Timer(true); // true makes this a daemon thread
+    pollTimer.schedule(
+        new TimerTask() {
+          @Override
+          public void run() {
+            updateBaggageRestrictions();
+          }
+        },
+        initialDelayMs,
+        refreshIntervalMs);
+  }
+
+  public boolean isReady() {
+    return initialized;
+  }
+
+  void updateBaggageRestrictions() {
+    List<BaggageRestrictionResponse> response;
+    try {
+      response = proxy.getBaggageRestrictions(serviceName);
+    } catch (BaggageRestrictionManagerException e) {
+      metrics.baggageRestrictionsUpdateFailure.inc(1);
+      return;
+    }
+
+    updateBaggageRestrictions(response);
+    metrics.baggageRestrictionsUpdateSuccess.inc(1);
+  }
+
+  private void updateBaggageRestrictions(List<BaggageRestrictionResponse> restrictions) {
+    Map<String, Restriction> baggageRestrictions = new HashMap<String, Restriction>();
+    for (BaggageRestrictionResponse restriction : restrictions) {
+      baggageRestrictions.put(restriction.getBaggageKey(), Restriction.of(true, restriction.getMaxValueLength()));
+    }
+    this.restrictions = baggageRestrictions;
+    initialized = true;
+  }
+
+  public void close() {
+    pollTimer.cancel();
+  }
+
+  @Override
+  public Restriction getRestriction(String service, String key) {
+    if (!initialized) {
+      if (denyBaggageOnInitializationFailure) {
+        return invalidRestriction;
+      } else {
+        return validRestriction;
+      }
+    }
+    Restriction restriction = this.restrictions.get(key);
+    if (restriction != null) {
+      return restriction;
+    }
+    return invalidRestriction;
+  }
+}

jaeger-core/src/main/java/com/uber/jaeger/baggage/http/BaggageRestrictionResponse.java

@@ -0,0 +1,31 @@
+/*
+ * Copyright (c) 2017, Uber Technologies, Inc
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+
+package com.uber.jaeger.baggage.http;
+
+import lombok.Value;
+
+@Value
+public class BaggageRestrictionResponse {
+  String baggageKey;
+  int maxValueLength;
+}

jaeger-core/src/main/java/com/uber/jaeger/exceptions/BaggageRestrictionManagerException.java

@@ -0,0 +1,35 @@
+/*
+ * Copyright (c) 2017, Uber Technologies, Inc
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+
+package com.uber.jaeger.exceptions;
+
+import java.io.IOException;
+
+public class BaggageRestrictionManagerException extends IOException {
+  public BaggageRestrictionManagerException(String msg) {
+    super(msg);
+  }
+
+  public BaggageRestrictionManagerException(String msg, Throwable cause) {
+    super(msg, cause);
+  }
+}

jaeger-core/src/main/java/com/uber/jaeger/metrics/Metrics.java

@@ -229,4 +229,18 @@ public static Metrics fromStatsReporter(StatsReporter reporter) {
   @Metric(name = "baggage-truncate")
   // Number of times baggage was truncated as per baggage restrictions
   public Counter baggageTruncate;
+
+  @Metric(
+      name = "baggage-restrictions-update",
+      tags = {@Tag(key = "result", value = "ok")}
+  )
+  // Number of times baggage restrictions were successfully updated
+  public Counter baggageRestrictionsUpdateSuccess;
+
+  @Metric(
+      name = "baggage-restrictions-update",
+      tags = {@Tag(key = "result", value = "err")}
+  )
+  // Number of times baggage restrictions failed to update
+  public Counter baggageRestrictionsUpdateFailure;
 }