brave.baggage.CorrelationUpdateScope Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of brave Show documentation
Brave: Core Libary
There is a newer version: 6.3.0
/*
 * Copyright The OpenZipkin Authors
 * SPDX-License-Identifier: Apache-2.0
 */
package brave.baggage;

import brave.baggage.CorrelationScopeConfig.SingleCorrelationField;
import brave.internal.CorrelationContext;
import brave.internal.Nullable;
import brave.propagation.CurrentTraceContext.Scope;
import java.util.concurrent.atomic.AtomicBoolean;

import static brave.baggage.CorrelationScopeDecorator.equal;
import static brave.baggage.CorrelationScopeDecorator.isSet;
import static brave.baggage.CorrelationScopeDecorator.setBit;

/** Handles reverting potentially late value updates to baggage fields. */
abstract class CorrelationUpdateScope extends AtomicBoolean implements Scope {
  CorrelationContext context;

  CorrelationUpdateScope(CorrelationContext context) {
    this.context = context;
  }

  /**
   * Called to get the name of the field, before it is flushed to the underlying context.
   *
   * @return the name used for this field, or null if not configured.
   */
  @Nullable abstract String name(BaggageField field);

  /**
   * Called after a field value is flushed to the underlying context. Only take action if the input
   * field is current being tracked.
   */
  abstract void handleUpdate(BaggageField field, @Nullable String value);

  static final class Single extends CorrelationUpdateScope {
    final Scope delegate;
    final SingleCorrelationField field;
    final @Nullable String valueToRevert;
    boolean shouldRevert;

    Single(
      Scope delegate,
      CorrelationContext context,
      SingleCorrelationField field,
      @Nullable String valueToRevert,
      boolean shouldRevert
    ) {
      super(context);
      this.delegate = delegate;
      this.field = field;
      this.valueToRevert = valueToRevert;
      this.shouldRevert = shouldRevert;
    }

    @Override public void close() {
      // don't duplicate work if called multiple times.
      if (!compareAndSet(false, true)) return;
      delegate.close();
      if (shouldRevert) context.update(field.name, valueToRevert);
    }

    @Override String name(BaggageField field) {
      return field.name;
    }

    @Override void handleUpdate(BaggageField field, String value) {
      if (!this.field.baggageField.equals(field)) return;
      if (!equal(value, valueToRevert)) shouldRevert = true;
    }
  }

  static final class Multiple extends CorrelationUpdateScope {
    final Scope delegate;
    final SingleCorrelationField[] fields;
    final String[] valuesToRevert;
    int shouldRevert;

    Multiple(
      Scope delegate,
      CorrelationContext context,
      SingleCorrelationField[] fields,
      String[] valuesToRevert,
      int shouldRevert
    ) {
      super(context);
      this.delegate = delegate;
      this.fields = fields;
      this.valuesToRevert = valuesToRevert;
      this.shouldRevert = shouldRevert;
    }

    @Override public void close() {
      // don't duplicate work if called multiple times.
      if (!compareAndSet(false, true)) return;

      delegate.close();
      for (int i = 0; i < fields.length; i++) {
        if (isSet(shouldRevert, i)) context.update(fields[i].name, valuesToRevert[i]);
      }
    }

    @Override String name(BaggageField field) {
      for (int i = 0; i < fields.length; i++) {
        if (fields[i].baggageField.equals(field)) {
          return fields[i].name;
        }
      }
      return null;
    }

    @Override void handleUpdate(BaggageField field, String value) {
      for (int i = 0; i < fields.length; i++) {
        if (fields[i].baggageField.equals(field)) {
          if (!equal(value, valuesToRevert[i])) shouldRevert = setBit(shouldRevert, i);
          return;
        }
      }
    }
  }
}