[TEX] Part 7: wiring data collection

atlasdb-client/src/main/java/com/palantir/atlasdb/transaction/impl/AbstractTransactionManager.java

@@ -81,6 +81,7 @@ protected final <T, E extends Exception> T runTaskThrowOnConflictWithCallback(
             }
         } finally {
             callback.run();
+            txn.reportExpectationsCollectedData();
             txn.runSuccessCallbacksIfDefinitivelyCommitted();
         }
     }

atlasdb-client/src/main/java/com/palantir/atlasdb/transaction/impl/CallbackAwareTransaction.java

@@ -17,7 +17,6 @@
 package com.palantir.atlasdb.transaction.impl;
 
 import com.palantir.atlasdb.transaction.api.PreCommitCondition;
-import com.palantir.atlasdb.transaction.api.Transaction;
 import com.palantir.atlasdb.transaction.api.TransactionFailedException;
 import com.palantir.atlasdb.transaction.service.TransactionService;
 
@@ -26,7 +25,7 @@
  * are run. This can be used to e.g run {@link PreCommitCondition#cleanup()} <i>before</i> the onSuccess callbacks
  * are run.
  */
-public interface CallbackAwareTransaction extends Transaction {
+public interface CallbackAwareTransaction extends ExpectationsAwareTransaction {
     void commitWithoutCallbacks() throws TransactionFailedException;
 
     void commitWithoutCallbacks(TransactionService transactionService) throws TransactionFailedException;

atlasdb-api/src/main/java/com/palantir/atlasdb/transaction/api/expectations/ExpectationsAwareTransaction.java → atlasdb-client/src/main/java/com/palantir/atlasdb/transaction/impl/ExpectationsAwareTransaction.java

@@ -14,15 +14,15 @@
  * limitations under the License.
  */
 
-package com.palantir.atlasdb.transaction.api.expectations;
+package com.palantir.atlasdb.transaction.impl;
 
 import com.palantir.atlasdb.transaction.api.Transaction;
+import com.palantir.atlasdb.transaction.api.expectations.TransactionReadInfo;
 
 /**
  * Implementors of this interface provide methods useful for tracking transactional expectations and whether
  * they were breached as well as relevant metrics and alerts. Transactional expectations represent transaction-level
  * limits and rules for proper usage of AtlasDB transactions (e.g. reading too much data overall).
- * Todo(aalouane): move this out of API once part 4 is merged
  */
 public interface ExpectationsAwareTransaction extends Transaction {
     long getAgeMillis();

atlasdb-client/src/main/java/com/palantir/atlasdb/transaction/impl/ForwardingCallbackAwareTransaction.java

@@ -19,7 +19,7 @@
 import com.palantir.atlasdb.transaction.api.TransactionFailedException;
 import com.palantir.atlasdb.transaction.service.TransactionService;
 
-public abstract class ForwardingCallbackAwareTransaction extends ForwardingTransaction
+public abstract class ForwardingCallbackAwareTransaction extends ForwardingExpectationsAwareTransaction
         implements CallbackAwareTransaction {
 
     @Override

atlasdb-client/src/main/java/com/palantir/atlasdb/transaction/impl/ForwardingExpectationsAwareTransaction.java

@@ -0,0 +1,41 @@
+/*
+ * (c) Copyright 2022 Palantir Technologies Inc. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.palantir.atlasdb.transaction.impl;
+
+import com.palantir.atlasdb.transaction.api.expectations.TransactionReadInfo;
+
+public abstract class ForwardingExpectationsAwareTransaction extends ForwardingTransaction
+        implements ExpectationsAwareTransaction {
+
+    @Override
+    public abstract ExpectationsAwareTransaction delegate();
+
+    @Override
+    public long getAgeMillis() {
+        return delegate().getAgeMillis();
+    }
+
+    @Override
+    public TransactionReadInfo getReadInfo() {
+        return delegate().getReadInfo();
+    }
+
+    @Override
+    public void reportExpectationsCollectedData() {
+        delegate().reportExpectationsCollectedData();
+    }
+}

atlasdb-impl-shared/build.gradle

@@ -12,6 +12,7 @@ license {
   exclude '**/PutUnlessExistsTableMetrics.java'
   exclude '**/TargetedSweepProgressMetrics.java'
   exclude '**/TimestampCorrectnessMetrics.java'
+  exclude '**/ExpectationsDataCollectionMetrics.java'
   exclude '**/VerificationModeMetrics.java'
 }
 

atlasdb-impl-shared/src/main/java/com/palantir/atlasdb/transaction/impl/SnapshotTransaction.java

@@ -96,6 +96,10 @@
 import com.palantir.atlasdb.transaction.api.TransactionLockAcquisitionTimeoutException;
 import com.palantir.atlasdb.transaction.api.TransactionLockTimeoutException;
 import com.palantir.atlasdb.transaction.api.TransactionReadSentinelBehavior;
+import com.palantir.atlasdb.transaction.api.expectations.TransactionReadInfo;
+import com.palantir.atlasdb.transaction.expectations.ExpectationsDataCollectionMetrics;
+import com.palantir.atlasdb.transaction.impl.expectations.TrackingKeyValueService;
+import com.palantir.atlasdb.transaction.impl.expectations.TrackingKeyValueServiceImpl;
 import com.palantir.atlasdb.transaction.impl.metrics.TableLevelMetricsController;
 import com.palantir.atlasdb.transaction.impl.metrics.TransactionOutcomeMetrics;
 import com.palantir.atlasdb.transaction.knowledge.TransactionKnowledgeComponents;
@@ -216,7 +220,7 @@ private enum State {
 
     protected final TimelockService timelockService;
     protected final LockWatchManagerInternal lockWatchManager;
-    final KeyValueService keyValueService;
+    final TrackingKeyValueService keyValueService;
     final AsyncKeyValueService immediateKeyValueService;
     final TransactionService defaultTransactionService;
     private final AsyncTransactionService immediateTransactionService;
@@ -266,14 +270,16 @@ private enum State {
 
     protected final TransactionKnowledgeComponents knowledge;
 
+    private final ExpectationsDataCollectionMetrics expectationsDataCollectionMetrics;
+
     /**
      * @param immutableTimestamp If we find a row written before the immutableTimestamp we don't need to grab a read
      *                           lock for it because we know that no writers exist.
      * @param preCommitCondition This check must pass for this transaction to commit.
      */
     /* package */ SnapshotTransaction(
             MetricsManager metricsManager,
-            KeyValueService keyValueService,
+            KeyValueService delegateKeyValueService,
             TimelockService timelockService,
             LockWatchManagerInternal lockWatchManager,
             TransactionService transactionService,
@@ -302,7 +308,7 @@ private enum State {
         this.lockWatchManager = lockWatchManager;
         this.conflictTracer = conflictTracer;
         this.transactionTimerContext = getTimer("transactionMillis").time();
-        this.keyValueService = keyValueService;
+        this.keyValueService = new TrackingKeyValueServiceImpl(delegateKeyValueService);
         this.immediateKeyValueService = KeyValueServices.synchronousAsAsyncKeyValueService(keyValueService);
         this.timelockService = timelockService;
         this.defaultTransactionService = transactionService;
@@ -338,6 +344,8 @@ private enum State {
                 timelockService,
                 immutableTimestamp,
                 knowledge);
+        this.expectationsDataCollectionMetrics =
+                ExpectationsDataCollectionMetrics.of(metricsManager.getTaggedRegistry());
     }
 
     protected TransactionScopedCache getCache() {
@@ -2576,6 +2584,34 @@ private Multimap<TableReference, Cell> getCellsToScrubByTable(State expectedStat
         return tableRefToCells;
     }
 
+    @Override
+    public long getAgeMillis() {
+        return System.currentTimeMillis() - timeCreated;
+    }
+
+    @Override
+    public TransactionReadInfo getReadInfo() {
+        return keyValueService.getOverallReadInfo();
+    }
+
+    @Override
+    public void reportExpectationsCollectedData() {
+        if (!(isDefinitivelyCommitted() || isAborted())) {
+            log.error("reportExpectationsCollectedData is called on an in-progress transaction");
+            return;
+        }
+
+        expectationsDataCollectionMetrics.ageMillis().update(getAgeMillis());
+
+        TransactionReadInfo info = getReadInfo();
+        expectationsDataCollectionMetrics.bytesRead().update(info.bytesRead());
+        expectationsDataCollectionMetrics.kvsCalls().update(info.kvsCalls());
+
+        info.maximumBytesKvsCallInfo()
+                .ifPresent(kvsCallReadInfo ->
+                        expectationsDataCollectionMetrics.worstKvsBytesRead().update(kvsCallReadInfo.bytesRead()));
+    }
+
     private Timer getTimer(String name) {
         return metricsManager.registerOrGetTimer(SnapshotTransaction.class, name);
     }

atlasdb-impl-shared/src/main/java/com/palantir/atlasdb/transaction/impl/SnapshotTransactionManager.java

@@ -183,6 +183,7 @@ public <T, C extends PreCommitCondition, E extends Exception> T runTaskWithCondi
             condition.cleanup();
             throw e;
         }
+
         return openTransaction.finishWithCallback(
                 transaction -> task.execute(transaction, condition), condition::cleanup);
     }
@@ -246,7 +247,7 @@ private OpenTransactionImpl(CallbackAwareTransaction delegate, LockToken immutab
         }
 
         @Override
-        public Transaction delegate() {
+        public CallbackAwareTransaction delegate() {
             return delegate;
         }
 

atlasdb-impl-shared/src/main/metrics/expectations.yml

@@ -0,0 +1,19 @@
+options:
+  javaPackage: 'com.palantir.atlasdb.transaction.expectations'
+
+namespaces:
+  expectationsDataCollection:
+    docs: Data collection on transaction patterns
+    metrics:
+      bytesRead:
+        docs: Bytes read by a transaction through Atlas kvs read calls
+        type: histogram
+      kvsCalls:
+        docs: Number of Atlas kvs read calls made by a transaction
+        type: histogram
+      ageMillis:
+        docs: Elapsed duration in milliseconds for a transaction
+        type: histogram
+      worstKvsBytesRead:
+        docs: Most bytes read by a transaction in one Atlas kvs read call
+        type: histogram