[TEX] Wrap KVS service reads and track bytes read (no metrics)

.palantir/revapi.yml

@@ -110,3 +110,8 @@ acceptedBreaks:
     - code: "java.method.removed"
       old: "method boolean com.palantir.atlasdb.transaction.service.TransactionStatus::equals(java.lang.Object)"
       justification: "Internal methods"
+  "0.815.0":
+    com.palantir.atlasdb:atlasdb-api:
+    - code: "java.class.removed"
+      old: "interface com.palantir.atlasdb.transaction.api.expectations.ExpectationsAwareTransaction"
+      justification: "moving out of API"

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

@@ -1,5 +1,5 @@
 /*
- * (c) Copyright 2022 Palantir Technologies Inc. All rights reserved.
+ * (c) Copyright 2023 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.
@@ -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 2023 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/src/main/java/com/palantir/atlasdb/transaction/impl/SnapshotTransaction.java

@@ -97,6 +97,9 @@
 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.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;
@@ -219,7 +222,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;
@@ -277,7 +280,7 @@ private enum State {
      */
     /* package */ SnapshotTransaction(
             MetricsManager metricsManager,
-            KeyValueService keyValueService,
+            KeyValueService delegateKeyValueService,
             TimelockService timelockService,
             LockWatchManagerInternal lockWatchManager,
             TransactionService transactionService,
@@ -306,7 +309,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;
@@ -2583,6 +2586,20 @@ private boolean validationNecessaryForInvolvedTablesOnCommit() {
         return anyTableRequiresImmutableTimestampLocking && needsToValidate;
     }
 
+    @Override
+    public long getAgeMillis() {
+        return System.currentTimeMillis() - timeCreated;
+    }
+
+    @Override
+    public TransactionReadInfo getReadInfo() {
+        return keyValueService.getOverallReadInfo();
+    }
+
+    // TODO(aalouane): implement and call this on transaction exit
+    @Override
+    public void reportExpectationsCollectedData() {}
+
     private long getStartTimestamp() {
         return startTimestamp.get();
     }

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

@@ -246,7 +246,7 @@ private OpenTransactionImpl(CallbackAwareTransaction delegate, LockToken immutab
         }
 
         @Override
-        public Transaction delegate() {
+        public CallbackAwareTransaction delegate() {
             return delegate;
         }
 

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

@@ -0,0 +1,5 @@
+type: feature
+feature:
+  description: '[TEX] Wrap KVS service reads and track bytes read (no metrics)'
+  links:
+  - https://github.com/palantir/atlasdb/pull/6475