Timelock uses tritium-preprocessor instrumentation (#7132)

Timelock uses tritium-preprocessor instrumentation

Author: schlosna

atlasdb-client/src/main/java/com/palantir/atlasdb/util/AtlasDbMetrics.java

@@ -23,6 +23,7 @@
 import com.palantir.tritium.api.event.InstrumentationFilter;
 import com.palantir.tritium.event.InstrumentationFilters;
 import com.palantir.tritium.event.InvocationContext;
+import com.palantir.tritium.event.InvocationEventHandler;
 import com.palantir.tritium.metrics.caffeine.CaffeineCacheStats;
 import com.palantir.tritium.metrics.registry.TaggedMetricRegistry;
 import com.palantir.tritium.proxy.Instrumentation;
@@ -81,11 +82,18 @@ public static <T, U extends T> T instrumentWithTaggedMetrics(
             U service,
             Function<InvocationContext, Map<String, String>> tagFunction) {
         return Instrumentation.builder(serviceInterface, service)
-                .withHandler(new TaggedMetricsInvocationEventHandler(
-                        taggedMetrics, MetricRegistry.name(serviceInterface), tagFunction))
+                .withHandler(taggedMetricsHandler(taggedMetrics, serviceInterface, tagFunction))
                 .build();
     }
 
+    public static <T> InvocationEventHandler<InvocationContext> taggedMetricsHandler(
+            TaggedMetricRegistry taggedMetrics,
+            Class<T> serviceInterface,
+            Function<InvocationContext, Map<String, String>> tagFunction) {
+        return new TaggedMetricsInvocationEventHandler(
+                taggedMetrics, MetricRegistry.name(serviceInterface), tagFunction);
+    }
+
     public static void registerCache(MetricRegistry metricRegistry, Cache<?, ?> cache, String metricsPrefix) {
         Set<String> existingMetrics = metricRegistry.getMetrics().keySet().stream()
                 .filter(name -> name.startsWith(metricsPrefix))

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

@@ -0,0 +1,5 @@
+type: improvement
+improvement:
+  description: Timelock uses tritium-preprocessor instrumentation
+  links:
+  - https://github.com/palantir/atlasdb/pull/7132

lock-api/build.gradle

@@ -27,12 +27,15 @@ dependencies {
     implementation 'com.palantir.safe-logging:safe-logging'
     implementation 'com.palantir.safe-logging:preconditions'
     implementation 'com.palantir.refreshable:refreshable'
+    implementation 'com.palantir.tritium:tritium-api'
+    implementation 'com.palantir.tritium:tritium-core'
     implementation 'com.palantir.tritium:tritium-ids'
 
     implementation 'com.google.guava:guava'
     implementation 'com.palantir.common:streams'
     implementation 'com.palantir.conjure.java:conjure-lib'
     implementation 'com.palantir.tokens:auth-tokens'
+    implementation 'com.palantir.tritium:tritium-annotations'
     implementation 'io.dropwizard.metrics:metrics-core'
     implementation project(':commons-annotations')
     implementation project(':commons-executors')
@@ -57,6 +60,7 @@ dependencies {
     compileOnly project(":atlasdb-processors")
     annotationProcessor 'org.immutables:value'
     compileOnly 'org.immutables:value::annotations'
+    annotationProcessor 'com.palantir.tritium:tritium-processor'
 
     testAnnotationProcessor 'org.immutables:value'
     testCompileOnly 'org.immutables:value::annotations'

lock-api/src/main/java/com/palantir/lock/LockService.java

@@ -21,6 +21,7 @@
 import com.palantir.common.annotation.NonIdempotent;
 import com.palantir.logsafe.Safe;
 import com.palantir.processors.AutoDelegate;
+import com.palantir.tritium.annotations.Instrument;
 import java.math.BigInteger;
 import java.util.Set;
 import javax.annotation.Nullable;
@@ -39,6 +40,7 @@
 @Path("/lock")
 @AutoDelegate
 @Beta
+@Instrument
 public interface LockService extends RemoteLockService {
     /**
      * Attempts to acquire the requested set of locks. The locks are

timelock-agent/build.gradle

@@ -26,6 +26,7 @@ dependencies {
     implementation 'com.palantir.safe-logging:safe-logging'
     implementation 'com.palantir.sls.versions:sls-versions'
     implementation 'com.palantir.tokens:auth-tokens'
+    implementation 'com.palantir.tritium:tritium-annotations'
     implementation 'com.palantir.tritium:tritium-api'
     implementation 'com.palantir.tritium:tritium-registry'
     implementation 'com.zaxxer:HikariCP'

timelock-agent/src/main/java/com/palantir/atlasdb/timelock/paxos/LeadershipComponents.java

@@ -18,30 +18,41 @@
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.io.Closer;
+import com.palantir.atlasdb.timelock.AsyncTimelockService;
+import com.palantir.atlasdb.timelock.InstrumentedAsyncTimelockService;
+import com.palantir.atlasdb.timelock.TimelockNamespaces;
+import com.palantir.atlasdb.util.AtlasDbMetrics;
+import com.palantir.common.concurrent.ConcurrentMaps;
 import com.palantir.leader.LeaderElectionService;
 import com.palantir.leader.NotCurrentLeaderException;
 import com.palantir.leader.proxy.AwaitingLeadershipProxy;
 import com.palantir.leader.proxy.LeadershipCoordinator;
+import com.palantir.lock.InstrumentedLockService;
+import com.palantir.lock.LockService;
 import com.palantir.logsafe.logger.SafeLogger;
 import com.palantir.logsafe.logger.SafeLoggerFactory;
 import com.palantir.paxos.Client;
 import com.palantir.timelock.paxos.HealthCheckPinger;
 import com.palantir.timelock.paxos.LeaderPingHealthCheck;
 import com.palantir.timelock.paxos.NamespaceTracker;
+import com.palantir.tritium.event.InvocationContext;
+import com.palantir.tritium.metrics.registry.TaggedMetricRegistry;
 import java.io.Closeable;
 import java.io.IOException;
 import java.util.Collection;
 import java.util.List;
-import java.util.concurrent.ConcurrentHashMap;
+import java.util.Map;
 import java.util.concurrent.ConcurrentMap;
+import java.util.function.Function;
 import java.util.function.Supplier;
 import org.immutables.value.Value;
 
 public class LeadershipComponents {
 
     private static final SafeLogger log = SafeLoggerFactory.get(LeadershipComponents.class);
 
-    private final ConcurrentMap<Client, LeadershipContext> leadershipContextByClient = new ConcurrentHashMap<>();
+    private final ConcurrentMap<Client, LeadershipContext> leadershipContextByClient =
+            ConcurrentMaps.newWithExpectedEntries(TimelockNamespaces.estimatedClients());
     private final ShutdownAwareCloser closer = new ShutdownAwareCloser();
 
     private final NetworkClientFactories.Factory<LeadershipContext> leadershipContextFactory;
@@ -54,15 +65,42 @@ public class LeadershipComponents {
         this.healthCheckPingers = healthCheckPingers;
     }
 
-    public <T> T wrapInLeadershipProxy(Client client, Class<T> clazz, Supplier<T> delegateSupplier) {
+    public LeadershipProxies createServices(
+            Client client,
+            Supplier<AsyncTimelockService> asyncTimelockServiceSupplier,
+            Supplier<LockService> lockServiceSupplier) {
         LeadershipContext context = getOrCreateNewLeadershipContext(client);
-        T instance = AwaitingLeadershipProxy.newProxyInstance(clazz, delegateSupplier, context.leadershipCoordinator());
+        TaggedMetricRegistry metrics = context.leadershipMetrics().taggedMetrics();
+        Function<InvocationContext, Map<String, String>> tagger =
+                context.leadershipMetrics().suspectedLeaderTag();
+        return ImmutableLeadershipProxies.of(
+                InstrumentedAsyncTimelockService.builder(wrapInLeadershipProxy(
+                                AsyncTimelockService.class,
+                                asyncTimelockServiceSupplier,
+                                context.leadershipCoordinator()))
+                        .withHandler(AtlasDbMetrics.taggedMetricsHandler(metrics, AsyncTimelockService.class, tagger))
+                        .build(),
+                InstrumentedLockService.builder(wrapInLeadershipProxy(
+                                LockService.class, lockServiceSupplier, context.leadershipCoordinator()))
+                        .withHandler(AtlasDbMetrics.taggedMetricsHandler(metrics, LockService.class, tagger))
+                        .build());
+    }
 
+    private <T> T wrapInLeadershipProxy(
+            Class<T> clazz, Supplier<T> delegateSupplier, LeadershipCoordinator leadershipCoordinator) {
+        T instance = AwaitingLeadershipProxy.newProxyInstance(clazz, delegateSupplier, leadershipCoordinator);
         // this is acceptable since the proxy returned implements Closeable and needs to be closed
-        Closeable closeableInstance = (Closeable) instance;
-        closer.register(closeableInstance);
+        closer.register((Closeable) instance);
+        return instance;
+    }
+
+    @Value.Immutable(builder = false)
+    public interface LeadershipProxies {
+        @Value.Parameter
+        AsyncTimelockService asyncTimelockService();
 
-        return context.leadershipMetrics().instrument(clazz, instance);
+        @Value.Parameter
+        LockService lockService();
     }
 
     public void shutdown() {
@@ -143,7 +181,7 @@ synchronized void shutdown() {
     }
 
     @Value.Immutable
-    abstract static class LeadershipContext {
+    public abstract static class LeadershipContext {
         abstract LeaderElectionService leaderElectionService();
 
         abstract LeadershipCoordinator leadershipCoordinator();

timelock-agent/src/main/java/com/palantir/atlasdb/timelock/paxos/LeadershipContextFactory.java

@@ -139,15 +139,16 @@ public LeadershipContext create(Client client) {
         BatchingLeaderElectionService leaderElectionService =
                 leaderElectionServiceFactory().create(clientAwareComponents);
 
-        LeadershipCoordinator leadershipCoordinator =
-                leadershipCoordinatorFactory().create(leaderElectionService);
+        LeadershipCoordinatorFactory coordinatorFactory = leadershipCoordinatorFactory();
+        LeadershipCoordinator leadershipCoordinator = coordinatorFactory.create(leaderElectionService);
 
         return ImmutableLeadershipContext.builder()
                 .leadershipMetrics(clientAwareComponents.leadershipMetrics())
                 .leaderElectionService(leaderElectionService)
                 .addCloseables(leaderElectionService)
                 .leadershipCoordinator(leadershipCoordinator)
-                .addCloseables(leadershipCoordinatorFactory())
+                .addCloseables(leadershipCoordinator)
+                .addCloseables(coordinatorFactory)
                 .addAllCloseables(leaderPingerFactory().closeables())
                 .build();
     }

timelock-agent/src/main/java/com/palantir/atlasdb/timelock/paxos/TimelockLeadershipMetrics.java

@@ -21,22 +21,25 @@
 import com.google.common.collect.SetMultimap;
 import com.palantir.atlasdb.AtlasDbMetricNames;
 import com.palantir.atlasdb.timelock.paxos.AutobatchingLeadershipObserverFactory.LeadershipEvent;
-import com.palantir.atlasdb.util.AtlasDbMetrics;
 import com.palantir.leader.LeadershipObserver;
 import com.palantir.leader.PaxosLeadershipEventRecorder;
 import com.palantir.logsafe.SafeArg;
 import com.palantir.logsafe.exceptions.SafeIllegalArgumentException;
 import com.palantir.paxos.Client;
+import com.palantir.tritium.event.InvocationContext;
 import com.palantir.tritium.metrics.registry.MetricName;
+import com.palantir.tritium.metrics.registry.TaggedMetricRegistry;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.function.Consumer;
+import java.util.function.Function;
 import java.util.function.Predicate;
 import org.immutables.value.Value;
 
 @Value.Immutable
 public abstract class TimelockLeadershipMetrics implements Dependencies.LeadershipMetrics {
-    @Value.Derived
+    @Value.Lazy
     List<SafeArg<String>> namespaceAsLoggingArgs() {
         return ImmutableList.of(
                 SafeArg.of(
@@ -45,28 +48,30 @@ List<SafeArg<String>> namespaceAsLoggingArgs() {
                 SafeArg.of(AtlasDbMetricNames.TAG_CLIENT, proxyClient().value()));
     }
 
-    @Value.Derived
+    @Value.Lazy
     public PaxosLeadershipEventRecorder eventRecorder() {
         return PaxosLeadershipEventRecorder.create(
-                metrics().clientScopedMetrics().metricRegistryForClient(proxyClient()), // metrics for client etc.
+                taggedMetrics(), // metrics for client etc.
                 leaderUuid().toString(),
                 leadershipObserver(),
                 namespaceAsLoggingArgs());
     }
 
-    @Value.Derived
+    @Value.Lazy
     LeadershipObserver leadershipObserver() {
         return leadershipObserverFactory().create(proxyClient());
     }
 
-    public <T> T instrument(Class<T> clazz, T instance) {
-        return AtlasDbMetrics.instrumentWithTaggedMetrics(
-                metrics().clientScopedMetrics().metricRegistryForClient(proxyClient()),
-                clazz,
-                instance,
-                _context -> ImmutableMap.of(
-                        AtlasDbMetricNames.TAG_CURRENT_SUSPECTED_LEADER,
-                        String.valueOf(localPingableLeader().ping())));
+    @Value.Lazy
+    public Function<InvocationContext, Map<String, String>> suspectedLeaderTag() {
+        return _context -> ImmutableMap.of(
+                AtlasDbMetricNames.TAG_CURRENT_SUSPECTED_LEADER,
+                String.valueOf(localPingableLeader().ping()));
+    }
+
+    @Value.Lazy
+    public TaggedMetricRegistry taggedMetrics() {
+        return metrics().clientScopedMetrics().metricRegistryForClient(proxyClient());
     }
 
     public static AutobatchingLeadershipObserverFactory createFactory(TimelockPaxosMetrics metrics) {

timelock-agent/src/main/java/com/palantir/timelock/paxos/AsyncTimeLockServicesCreator.java

@@ -27,6 +27,7 @@
 import com.palantir.atlasdb.timelock.lockwatches.BufferMetrics;
 import com.palantir.atlasdb.timelock.lockwatches.RequestMetrics;
 import com.palantir.atlasdb.timelock.paxos.LeadershipComponents;
+import com.palantir.atlasdb.timelock.paxos.LeadershipComponents.LeadershipProxies;
 import com.palantir.atlasdb.util.MetricsManager;
 import com.palantir.common.concurrent.NamedThreadFactory;
 import com.palantir.common.concurrent.PTExecutors;
@@ -68,17 +69,12 @@ public TimeLockServices createTimeLockServices(
             Supplier<LockService> rawLockServiceSupplier) {
         log.info("Creating async timelock services for client {}", SafeArg.of("client", client));
         LockLog maybeEnhancedLockLog = maybeEnhancedLockLog(client);
-
-        AsyncTimelockService asyncTimelockService = leadershipComponents.wrapInLeadershipProxy(
-                client,
-                AsyncTimelockService.class,
-                () -> createRawAsyncTimelockService(client, rawTimestampServiceSupplier, maybeEnhancedLockLog));
-
-        LockService lockService = leadershipComponents.wrapInLeadershipProxy(
+        LeadershipProxies leadershipProxies = leadershipComponents.createServices(
                 client,
-                LockService.class,
+                () -> createRawAsyncTimelockService(client, rawTimestampServiceSupplier, maybeEnhancedLockLog),
                 Suppliers.compose(NonTransactionalLockService::new, rawLockServiceSupplier::get));
-
+        AsyncTimelockService asyncTimelockService = leadershipProxies.asyncTimelockService();
+        LockService lockService = leadershipProxies.lockService();
         return TimeLockServices.create(
                 asyncTimelockService, lockService, asyncTimelockService, asyncTimelockService, maybeEnhancedLockLog);
     }