Reuse a single TrackedDataInputPlus instance per UnfilteredSerializer

patch by Francisco Guerrero; reviewed Dmitry Konstantinov, Caleb Rackliffe for CASSANDRA-21296

CHANGES.txt

@@ -1,4 +1,5 @@
 6.0-alpha2
+ * Reuse a single TrackedDataInputPlus instance per UnfilteredSerializer (CASSANDRA-21296)
  * In-tree docs are included in binary artifacts (tarball, deb, rpm) as html and man pages (CASSANDRA-17260)
  * Add tooling to repair system peers and peers_v2 if inconsistent with cluster metadata (CASSANDRA-21187)
  * Fix a removed TTLed row re-appearance in a materialized view after a cursor compaction (CASSANDRA-21152)

src/java/org/apache/cassandra/db/rows/DeserializationHelper.java

@@ -17,6 +17,7 @@
  */
 package org.apache.cassandra.db.rows;
 
+import java.io.DataInput;
 import java.nio.ByteBuffer;
 import java.util.Map;
 
@@ -25,6 +26,7 @@ import org.apache.cassandra.db.LivenessInfo;
 import org.apache.cassandra.db.context.CounterContext;
 import org.apache.cassandra.db.filter.ColumnFilter;
 import org.apache.cassandra.db.marshal.ValueAccessor;
+import org.apache.cassandra.io.util.TrackedDataInputPlus;
 import org.apache.cassandra.schema.ColumnMetadata;
 import org.apache.cassandra.schema.DroppedColumn;
 import org.apache.cassandra.schema.TableMetadata;
@@ -57,6 +59,8 @@ public class DeserializationHelper
     private final Map<ByteBuffer, DroppedColumn> droppedColumns;
     private DroppedColumn currentDroppedComplex;
 
+    // Reusable per-partition tracker for row-size-bounded reads in SSTable deserialization.
+    private TrackedDataInputPlus trackedInput;
 
     public DeserializationHelper(TableMetadata metadata, int version, Flag flag, ColumnFilter columnsToFetch)
     {
@@ -148,4 +152,19 @@ public class DeserializationHelper
                ? CounterContext.instance().clearAllLocal(value, accessor)
                : value;
     }
+
+    /**
+     * @param source the original source of {@link DataInput}
+     * @param limit  the limit number of bytes to read
+     * @return a reusable {@link TrackedDataInputPlus}. The instance is lazily created on
+     * first use and reused for every row in the partition.
+     */
+    public TrackedDataInputPlus trackedDataInputPlus(DataInput source, long limit)
+    {
+        if (trackedInput == null)
+            trackedInput = new TrackedDataInputPlus(source, limit);
+        else
+            trackedInput.reset(source, limit);
+        return trackedInput;
+    }
 }

src/java/org/apache/cassandra/db/rows/UnfilteredSerializer.java

@@ -36,7 +36,6 @@ import org.apache.cassandra.io.util.DataInputPlus;
 import org.apache.cassandra.io.util.DataOutputBuffer;
 import org.apache.cassandra.io.util.DataOutputPlus;
 import org.apache.cassandra.io.util.FileDataInput;
-import org.apache.cassandra.io.util.TrackedDataInputPlus;
 import org.apache.cassandra.net.MessagingService;
 import org.apache.cassandra.schema.ColumnMetadata;
 import org.apache.cassandra.utils.SearchIterator;
@@ -601,7 +600,7 @@ public class UnfilteredSerializer
             {
                 long rowSize = in.readUnsignedVInt();
                 in.readUnsignedVInt(); // previous unfiltered size
-                in = new TrackedDataInputPlus(in, rowSize);
+                in = helper.trackedDataInputPlus(in, rowSize);
             }
 
             LivenessInfo rowLiveness = LivenessInfo.EMPTY;

src/java/org/apache/cassandra/io/util/TrackedDataInputPlus.java

@@ -22,6 +22,8 @@ import java.io.DataInputStream;
 import java.io.EOFException;
 import java.io.IOException;
 
+import javax.annotation.concurrent.NotThreadSafe;
+
 import net.nicoulaj.compilecommand.annotations.Inline;
 
 import org.apache.cassandra.db.TypeSizes;
@@ -29,11 +31,12 @@ import org.apache.cassandra.db.TypeSizes;
 /**
  * This class is to track bytes read from given DataInput
  */
+@NotThreadSafe
 public class TrackedDataInputPlus implements DataInputPlus, BytesReadTracker
 {
     private long bytesRead;
-    private final long limit;
-    final DataInput source;
+    private long limit;
+    private DataInput source;
 
     /**
      * Create a TrackedDataInputPlus from given DataInput with no limit of bytes to read
@@ -66,6 +69,19 @@ public class TrackedDataInputPlus implements DataInputPlus, BytesReadTracker
         bytesRead = count;
     }
 
+    /**
+     * Resets the tracked {@link DataInputPlus} with the given {@code source} and {@code limit}
+     *
+     * @param source the original source of {@link DataInput}
+     * @param limit  the limit number of bytes to read
+     */
+    public void reset(DataInput source, long limit)
+    {
+        this.source = source;
+        this.limit = limit;
+        this.bytesRead = 0;
+    }
+
     public boolean readBoolean() throws IOException
     {
         checkCanRead(TypeSizes.BOOL_SIZE);

test/unit/org/apache/cassandra/db/rows/DeserializationHelperTest.java

@@ -0,0 +1,51 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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 org.apache.cassandra.db.rows;
+
+
+import java.io.DataInput;
+
+import org.junit.Test;
+
+import org.apache.cassandra.db.marshal.AsciiType;
+import org.apache.cassandra.db.marshal.Int32Type;
+import org.apache.cassandra.io.util.TrackedDataInputPlus;
+import org.apache.cassandra.schema.TableMetadata;
+
+import static org.apache.cassandra.net.MessagingService.Version.VERSION_60;
+import static org.junit.Assert.assertSame;
+import static org.mockito.Mockito.mock;
+
+public class DeserializationHelperTest
+{
+    static TableMetadata metadata =
+    TableMetadata.builder("DeserializationHelperTest", "Test")
+                 .addPartitionKeyColumn("key", AsciiType.instance)
+                 .addClusteringColumn("clustering", Int32Type.instance)
+                 .addRegularColumn("data", Int32Type.instance)
+                 .build();
+
+    @Test
+    public void testTrackedDataInputPlusIsReusable()
+    {
+        DeserializationHelper helper = new DeserializationHelper(metadata, VERSION_60.value, DeserializationHelper.Flag.LOCAL);
+        TrackedDataInputPlus trackedDataInputPlus = helper.trackedDataInputPlus(mock(DataInput.class), 0);
+        assertSame(trackedDataInputPlus, helper.trackedDataInputPlus(mock(DataInput.class), 1));
+    }
+}