Add traces for kvs-get

atlasdb-cassandra/src/main/java/com/palantir/atlasdb/keyvalue/cassandra/CassandraKeyValueServiceImpl.java

@@ -97,6 +97,7 @@
 import com.palantir.common.streams.KeyedStream;
 import com.palantir.logsafe.Preconditions;
 import com.palantir.logsafe.SafeArg;
+import com.palantir.tracing.CloseableTracer;
 import com.palantir.tracing.Tracers;
 import com.palantir.tritium.metrics.MetricRegistries;
 import com.palantir.tritium.metrics.registry.TaggedMetricRegistry;
@@ -777,7 +778,9 @@ public Map<Cell, Value> get(TableReference tableRef, Map<Cell, Long> timestampBy
             ImmutableMap.Builder<Cell, Value> builder = ImmutableMap.builder();
             for (long ts : cellsByTs.keySet()) {
                 StartTsResultsCollector collector = new StartTsResultsCollector(metricsManager, ts);
-                cellLoader.loadWithTs("get", tableRef, cellsByTs.get(ts), ts, false, collector, readConsistency);
+                try (CloseableTracer tracer = CloseableTracer.startSpan("loadWithTs")) {
+                    cellLoader.loadWithTs("get", tableRef, cellsByTs.get(ts), ts, false, collector, readConsistency);
+                }
                 builder.putAll(collector.getCollectedResults());
             }
             return builder.build();
@@ -789,7 +792,10 @@ public Map<Cell, Value> get(TableReference tableRef, Map<Cell, Long> timestampBy
     private Map<Cell, Value> get(
             String kvsMethodName, TableReference tableRef, Set<Cell> cells, long maxTimestampExclusive) {
         StartTsResultsCollector collector = new StartTsResultsCollector(metricsManager, maxTimestampExclusive);
-        cellLoader.loadWithTs(kvsMethodName, tableRef, cells, maxTimestampExclusive, false, collector, readConsistency);
+        try (CloseableTracer tracer = CloseableTracer.startSpan("loadWithTs")) {
+            cellLoader.loadWithTs(
+                    kvsMethodName, tableRef, cells, maxTimestampExclusive, false, collector, readConsistency);
+        }
         return collector.getCollectedResults();
     }
 

atlasdb-cassandra/src/main/java/com/palantir/atlasdb/keyvalue/cassandra/CellLoader.java

@@ -15,6 +15,7 @@
  */
 package com.palantir.atlasdb.keyvalue.cassandra;
 
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Lists;
 import com.google.common.collect.Maps;
@@ -29,6 +30,7 @@
 import com.palantir.atlasdb.util.AnnotationType;
 import com.palantir.common.base.FunctionCheckedException;
 import com.palantir.logsafe.SafeArg;
+import com.palantir.tracing.CloseableTracer;
 import java.net.InetSocketAddress;
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
@@ -92,8 +94,20 @@ void loadWithTs(
             boolean loadAllTs,
             CassandraKeyValueServices.ThreadSafeResultVisitor visitor,
             ConsistencyLevel consistency) {
-        Map<InetSocketAddress, List<Cell>> hostsAndCells =
-                HostPartitioner.partitionByHost(clientPool, cells, Cell::getRowName);
+        Map<InetSocketAddress, List<Cell>> hostsAndCells;
+        try (CloseableTracer tracer = CloseableTracer.startSpan(
+                "partitionByHost",
+                ImmutableMap.of(
+                        "cells",
+                        String.valueOf(cells.size()),
+                        "tableRef",
+                        LoggingArgs.safeInternalTableNameOrPlaceholder(tableRef.toString()),
+                        "timestampClause",
+                        loadAllTs ? "for all timestamps " : "",
+                        "startTs",
+                        String.valueOf(startTs)))) {
+            hostsAndCells = HostPartitioner.partitionByHost(clientPool, cells, Cell::getRowName);
+        }
         int totalPartitions = hostsAndCells.keySet().size();
 
         if (log.isTraceEnabled()) {
@@ -118,15 +132,17 @@ void loadWithTs(
                         SafeArg.of("ipPort", hostAndCells.getKey()));
             }
 
-            tasks.addAll(getLoadWithTsTasksForSingleHost(
-                    kvsMethodName,
-                    hostAndCells.getKey(),
-                    tableRef,
-                    hostAndCells.getValue(),
-                    startTs,
-                    loadAllTs,
-                    visitor,
-                    consistency));
+            try (CloseableTracer tracer = CloseableTracer.startSpan("getLoadWithTsTasksForSingleHost")) {
+                tasks.addAll(getLoadWithTsTasksForSingleHost(
+                        kvsMethodName,
+                        hostAndCells.getKey(),
+                        tableRef,
+                        hostAndCells.getValue(),
+                        startTs,
+                        loadAllTs,
+                        visitor,
+                        consistency));
+            }
         }
 
         taskRunner.runAllTasksCancelOnFailure(tasks);

atlasdb-cassandra/src/main/java/com/palantir/atlasdb/keyvalue/cassandra/TaskRunner.java

@@ -15,19 +15,25 @@
  */
 package com.palantir.atlasdb.keyvalue.cassandra;
 
+import com.google.common.util.concurrent.FutureCallback;
+import com.google.common.util.concurrent.Futures;
+import com.google.common.util.concurrent.ListenableFuture;
+import com.google.common.util.concurrent.ListeningExecutorService;
+import com.google.common.util.concurrent.MoreExecutors;
 import com.palantir.common.base.Throwables;
+import com.palantir.tracing.DetachedSpan;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 import java.util.concurrent.Callable;
+import java.util.concurrent.Executor;
 import java.util.concurrent.ExecutorService;
-import java.util.concurrent.Future;
 
 class TaskRunner {
-    private ExecutorService executor;
+    private final ListeningExecutorService listeningExecutor;
 
     TaskRunner(ExecutorService executor) {
-        this.executor = executor;
+        this.listeningExecutor = MoreExecutors.listeningDecorator(executor);
     }
 
     /*
@@ -44,22 +50,43 @@ <V> List<V> runAllTasksCancelOnFailure(List<Callable<V>> tasks) {
             }
         }
 
-        List<Future<V>> futures = new ArrayList<>(tasks.size());
+        List<ListenableFuture<V>> futures = new ArrayList<>(tasks.size());
         for (Callable<V> task : tasks) {
-            futures.add(executor.submit(task));
+            DetachedSpan detachedSpan = DetachedSpan.start("task");
+            ListenableFuture<V> future = listeningExecutor.submit(task);
+            futures.add(attachDetachedSpanCompletion(detachedSpan, future, listeningExecutor));
         }
         try {
             List<V> results = new ArrayList<>(tasks.size());
-            for (Future<V> future : futures) {
+            for (ListenableFuture<V> future : futures) {
                 results.add(future.get());
             }
             return results;
         } catch (Exception e) {
             throw Throwables.unwrapAndThrowAtlasDbDependencyException(e);
         } finally {
-            for (Future<V> future : futures) {
+            for (ListenableFuture<V> future : futures) {
                 future.cancel(true);
             }
         }
     }
+
+    private static <V> ListenableFuture<V> attachDetachedSpanCompletion(
+            DetachedSpan detachedSpan, ListenableFuture<V> future, Executor tracingExecutorService) {
+        Futures.addCallback(
+                future,
+                new FutureCallback<V>() {
+                    @Override
+                    public void onSuccess(V result) {
+                        detachedSpan.complete();
+                    }
+
+                    @Override
+                    public void onFailure(Throwable throwable) {
+                        detachedSpan.complete();
+                    }
+                },
+                tracingExecutorService);
+        return future;
+    }
 }

changelog/@unreleased/pr-5466.v2.yml

@@ -0,0 +1,6 @@
+type: improvement
+improvement:
+  description: |-
+    Add traces inside kvs-get to get more granular performance information.
+  links:
+  - https://github.com/palantir/atlasdb/pull/5466