zipkin2.storage.cassandra.SelectTraceIdsFromServiceSpan Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of zipkin-storage-cassandra Show documentation
There is a newer version: 3.5.1
Show newest version
/*
 * Copyright The OpenZipkin Authors
 * SPDX-License-Identifier: Apache-2.0
 */
package zipkin2.storage.cassandra;

import com.datastax.oss.driver.api.core.CqlSession;
import com.datastax.oss.driver.api.core.cql.AsyncResultSet;
import com.datastax.oss.driver.api.core.cql.BoundStatementBuilder;
import com.datastax.oss.driver.api.core.cql.PreparedStatement;
import com.google.auto.value.AutoValue;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.CompletionStage;
import zipkin2.Call;
import zipkin2.internal.Nullable;
import zipkin2.storage.cassandra.CassandraSpanStore.TimestampRange;
import zipkin2.storage.cassandra.internal.call.AccumulateTraceIdTsUuid;
import zipkin2.storage.cassandra.internal.call.AggregateIntoMap;
import zipkin2.storage.cassandra.internal.call.ResultSetFutureCall;

import static zipkin2.storage.cassandra.Schema.TABLE_TRACE_BY_SERVICE_SPAN;

final class SelectTraceIdsFromServiceSpan extends ResultSetFutureCall {
  @AutoValue abstract static class Input {
    abstract String service();

    abstract String span();

    abstract int bucket();

    @Nullable abstract Long start_duration();

    @Nullable abstract Long end_duration();

    abstract UUID start_ts();

    abstract UUID end_ts();

    abstract int limit_();

    Input withService(String service) {
      return new AutoValue_SelectTraceIdsFromServiceSpan_Input(
        service,
        span(),
        bucket(),
        start_duration(),
        end_duration(),
        start_ts(),
        end_ts(),
        limit_());
    }
  }

  static final class Factory {
    final CqlSession session;
    final PreparedStatement selectTraceIdsByServiceSpanName;
    final PreparedStatement selectTraceIdsByServiceSpanNameAndDuration;

    Factory(CqlSession session) {
      this.session = session;
      String baseQuery = "SELECT trace_id,ts"
        + " FROM " + TABLE_TRACE_BY_SERVICE_SPAN
        + " WHERE service=?"
        + " AND span=?"
        + " AND bucket=?"
        + " AND ts>=?"
        + " AND ts<=?";
      this.selectTraceIdsByServiceSpanName = session.prepare(baseQuery
        + " LIMIT ?");
      this.selectTraceIdsByServiceSpanNameAndDuration = session.prepare(baseQuery
        + " AND duration>=?"
        + " AND duration<=?"
        + " LIMIT ?");
    }

    Input newInput(
      String serviceName,
      String spanName,
      int bucket,
      @Nullable Long minDurationMicros,
      @Nullable Long maxDurationMicros,
      TimestampRange timestampRange,
      int limit) {
      Long start_duration = null, end_duration = null;
      if (minDurationMicros != null) {
        start_duration = minDurationMicros / 1000L;
        end_duration = maxDurationMicros != null ? maxDurationMicros / 1000L : Long.MAX_VALUE;
      }
      return new AutoValue_SelectTraceIdsFromServiceSpan_Input(
        serviceName,
        spanName,
        bucket,
        start_duration,
        end_duration,
        timestampRange.startUUID,
        timestampRange.endUUID,
        limit);
    }

    Call> newCall(List inputs) {
      if (inputs.isEmpty()) return Call.create(Map.of());
      if (inputs.size() == 1) return newCall(inputs.get(0));

      List>> bucketedTraceIdCalls = new ArrayList<>();
      for (SelectTraceIdsFromServiceSpan.Input input : inputs) {
        bucketedTraceIdCalls.add(newCall(input));
      }
      return new AggregateIntoMap<>(bucketedTraceIdCalls);
    }

    Call> newCall(Input input) {
      PreparedStatement preparedStatement = input.start_duration() != null
        ? selectTraceIdsByServiceSpanNameAndDuration
        : selectTraceIdsByServiceSpanName;
      return new SelectTraceIdsFromServiceSpan(this, preparedStatement, input)
        .flatMap(AccumulateTraceIdTsUuid.get());
    }

    /** Applies all deferred service names to all input templates */
    FlatMapper, Map> newFlatMapper(List inputTemplates) {
      return new FlatMapServicesToInputs(inputTemplates);
    }

    class FlatMapServicesToInputs implements FlatMapper, Map> {
      final List inputTemplates;

      FlatMapServicesToInputs(List inputTemplates) {
        this.inputTemplates = inputTemplates;
      }

      @Override public Call> map(List serviceNames) {
        List>> bucketedTraceIdCalls = new ArrayList<>();

        for (String service : serviceNames) { // fan out every input for each service name
          List scopedInputs = new ArrayList<>();
          for (SelectTraceIdsFromServiceSpan.Input input : inputTemplates) {
            scopedInputs.add(input.withService(service));
          }
          bucketedTraceIdCalls.add(newCall(scopedInputs));
        }

        if (bucketedTraceIdCalls.isEmpty()) return Call.create(Map.of());
        if (bucketedTraceIdCalls.size() == 1) return bucketedTraceIdCalls.get(0);
        return new AggregateIntoMap<>(bucketedTraceIdCalls);
      }

      @Override public String toString() {
        List inputs = new ArrayList<>();
        for (Input input : inputTemplates) {
          inputs.add(input.toString().replace("Input", "SelectTraceIdsFromServiceSpan"));
        }
        return "FlatMapServicesToInputs{" + inputs + "}";
      }
    }
  }

  final Factory factory;
  final PreparedStatement preparedStatement;
  final Input input;

  SelectTraceIdsFromServiceSpan(Factory factory, PreparedStatement preparedStatement, Input input) {
    this.factory = factory;
    this.preparedStatement = preparedStatement;
    this.input = input;
  }

  @Override protected CompletionStage newCompletionStage() {
    int i = 0;
    BoundStatementBuilder bound = preparedStatement.boundStatementBuilder()
      .setString(i++, input.service())
      .setString(i++, input.span())
      .setInt(i++, input.bucket())
      .setUuid(i++, input.start_ts())
      .setUuid(i++, input.end_ts());

    if (input.start_duration() != null) {
      bound.setLong(i++, input.start_duration());
      bound.setLong(i++, input.end_duration());
    }

    bound
      .setInt(i, input.limit_())
      .setPageSize(input.limit_());

    return factory.session.executeAsync(bound.build());
  }

  @Override public AsyncResultSet map(AsyncResultSet input) {
    return input;
  }

  @Override public String toString() {
    return input.toString().replace("Input", "SelectTraceIdsFromServiceSpan");
  }

  @Override public SelectTraceIdsFromServiceSpan clone() {
    return new SelectTraceIdsFromServiceSpan(factory, preparedStatement, input);
  }
}