thread HFileConext using trace Context

Author: ndimiduk

hbase-common/src/main/java/org/apache/hadoop/hbase/io/encoding/HFileBlockDefaultDecodingContext.java

@@ -17,8 +17,6 @@
  */
 package org.apache.hadoop.hbase.io.encoding;
 
-import io.opentelemetry.api.common.Attributes;
-import io.opentelemetry.api.common.AttributesBuilder;
 import java.io.DataInputStream;
 import java.io.IOException;
 import java.io.InputStream;
@@ -32,7 +30,6 @@
 import org.apache.hadoop.hbase.io.crypto.Decryptor;
 import org.apache.hadoop.hbase.io.crypto.Encryption;
 import org.apache.hadoop.hbase.io.hfile.HFileContext;
-import org.apache.hadoop.hbase.io.hfile.trace.HFileContextAttributesBuilder;
 import org.apache.hadoop.hbase.io.util.BlockIOUtils;
 import org.apache.hadoop.hbase.nio.ByteBuff;
 import org.apache.hadoop.hbase.util.Bytes;
@@ -48,14 +45,11 @@
 public class HFileBlockDefaultDecodingContext implements HFileBlockDecodingContext {
   private final Configuration conf;
   private final HFileContext fileContext;
-  private final AttributesBuilder attributesBuilder;
   private TagCompressionContext tagCompressionContext;
 
   public HFileBlockDefaultDecodingContext(Configuration conf, HFileContext fileContext) {
     this.conf = conf;
     this.fileContext = fileContext;
-    this.attributesBuilder =
-      new HFileContextAttributesBuilder(Attributes.builder(), fileContext).build();
   }
 
   @Override
@@ -108,7 +102,7 @@ public void prepareDecoding(int onDiskSizeWithoutHeader, int uncompressedSizeWit
           try (InputStream is =
             compression.createDecompressionStream(dataInputStream, decompressor, 0)) {
             BlockIOUtils.readFullyWithHeapBuffer(is, blockBufferWithoutHeader,
-              uncompressedSizeWithoutHeader, attributesBuilder);
+              uncompressedSizeWithoutHeader);
           }
         } finally {
           if (decompressor != null) {
@@ -117,7 +111,7 @@ public void prepareDecoding(int onDiskSizeWithoutHeader, int uncompressedSizeWit
         }
       } else {
         BlockIOUtils.readFullyWithHeapBuffer(dataInputStream, blockBufferWithoutHeader,
-          onDiskSizeWithoutHeader, attributesBuilder);
+          onDiskSizeWithoutHeader);
       }
     } finally {
       byteBuffInputStream.close();

…trace/HFileContextAttributesBuilder.java …ileContextAttributesBuilderConsumer.javahbase-common/src/main/java/org/apache/hadoop/hbase/io/hfile/trace/HFileContextAttributesBuilder.java renamed to hbase-common/src/main/java/org/apache/hadoop/hbase/io/hfile/trace/HFileContextAttributesBuilderConsumer.java hbase-common/src/main/java/org/apache/hadoop/hbase/io/hfile/trace/HFileContextAttributesBuilder.java renamed to hbase-common/src/main/java/org/apache/hadoop/hbase/io/hfile/trace/HFileContextAttributesBuilderConsumer.java

@@ -22,45 +22,69 @@
 import static org.apache.hadoop.hbase.trace.HBaseSemanticAttributes.DATA_BLOCK_ENCODING_KEY;
 import static org.apache.hadoop.hbase.trace.HBaseSemanticAttributes.ENCRYPTION_CIPHER_KEY;
 import static org.apache.hadoop.hbase.trace.HBaseSemanticAttributes.HFILE_NAME_KEY;
+import static org.apache.hadoop.hbase.trace.HBaseSemanticAttributes.READ_TYPE_KEY;
 
 import io.opentelemetry.api.common.AttributesBuilder;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.ContextKey;
 import java.util.Objects;
-import java.util.function.Supplier;
+import java.util.function.Consumer;
 import org.apache.hadoop.hbase.io.hfile.HFileContext;
+import org.apache.hadoop.hbase.trace.HBaseSemanticAttributes.ReadType;
 import org.apache.hadoop.hbase.util.ChecksumType;
 import org.apache.yetus.audience.InterfaceAudience;
 
 /**
- * Populate fields on an {@link AttributesBuilder} based on an {@link HFileContext}.
+ * <p>
+ * Populate fields on an {@link AttributesBuilder} based on an {@link HFileContext}. Passed around
+ * inside an active {@link Context}, indexed under {@link #CONTEXT_KEY}. The class is designed such
+ * that calls to the {@link #accept(AttributesBuilder)} method are idempotent with regards to the
+ * instance of this class.
+ * </p>
+ * <p>
+ * The true and truly ridiculous class name should be something more like
+ * {@code HFileContext_ContextAttributes_AttributesBuilder_Consumer}.
+ * </p>
  */
 @InterfaceAudience.Private
-public class HFileContextAttributesBuilder implements Supplier<AttributesBuilder> {
+public class HFileContextAttributesBuilderConsumer implements Consumer<AttributesBuilder> {
+
+  /**
+   * Used to place extract attributes pertaining to the {@link HFileContext} that scopes the active
+   * {@link Context}.
+   */
+  public static final ContextKey<Consumer<AttributesBuilder>> CONTEXT_KEY =
+    ContextKey.named("db.hbase.io.hfile.context_attributes");
 
-  private final AttributesBuilder builder;
   private final HFileContext hFileContext;
 
   private boolean skipChecksum = false;
+  private ReadType readType = null;
 
-  public HFileContextAttributesBuilder(final AttributesBuilder builder,
-    final HFileContext hFileContext) {
-    this.builder = Objects.requireNonNull(builder);
+  public HFileContextAttributesBuilderConsumer(final HFileContext hFileContext) {
     this.hFileContext = Objects.requireNonNull(hFileContext);
   }
 
   /**
    * Specify that the {@link ChecksumType} should not be included in the attributes.
    */
-  public HFileContextAttributesBuilder setSkipChecksum(final boolean skipChecksum) {
+  public HFileContextAttributesBuilderConsumer setSkipChecksum(final boolean skipChecksum) {
     this.skipChecksum = skipChecksum;
     return this;
   }
 
-  @Override
-  public AttributesBuilder get() {
-    return build();
+  /**
+   * Specify the {@link ReadType} involced in this IO operation.
+   */
+  public HFileContextAttributesBuilderConsumer setReadType(final ReadType readType) {
+    // TODO: this is not a part of the HFileBlock, its context of the operation. Should track this
+    // detail elsewhere.
+    this.readType = readType;
+    return this;
   }
 
-  public AttributesBuilder build() {
+  @Override
+  public void accept(AttributesBuilder builder) {
     if (hFileContext.getHFileName() != null) {
       builder.put(HFILE_NAME_KEY, hFileContext.getHFileName());
     }
@@ -79,7 +103,8 @@ public AttributesBuilder build() {
     if (!skipChecksum && hFileContext.getChecksumType() != null) {
       builder.put(CHECKSUM_KEY, hFileContext.getChecksumType().getName());
     }
-
-    return builder;
+    if (readType != null) {
+      builder.put(READ_TYPE_KEY, readType.name());
+    }
   }
 }

hbase-common/src/main/java/org/apache/hadoop/hbase/io/util/BlockIOUtils.java

@@ -23,13 +23,16 @@
 import io.opentelemetry.api.common.Attributes;
 import io.opentelemetry.api.common.AttributesBuilder;
 import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.context.Context;
 import java.io.IOException;
 import java.io.InputStream;
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
 import java.nio.ByteBuffer;
+import java.util.Optional;
 import org.apache.hadoop.fs.ByteBufferReadable;
 import org.apache.hadoop.fs.FSDataInputStream;
+import org.apache.hadoop.hbase.io.hfile.trace.HFileContextAttributesBuilderConsumer;
 import org.apache.hadoop.hbase.nio.ByteBuff;
 import org.apache.hadoop.io.IOUtils;
 import org.apache.yetus.audience.InterfaceAudience;
@@ -81,17 +84,16 @@ public static boolean isByteBufferReadable(FSDataInputStream is) {
    * @param length bytes to read.
    * @throws IOException exception to throw if any error happen
    */
-  public static void readFully(ByteBuff buf, FSDataInputStream dis, int length,
-    AttributesBuilder attributesBuilder) throws IOException {
+  public static void readFully(ByteBuff buf, FSDataInputStream dis, int length) throws IOException {
     final Span span = Span.current();
-    final AttributesBuilder copy = Attributes.builder().putAll(attributesBuilder.build());
+    final AttributesBuilder attributesBuilder = builderFromContext(Context.current());
     if (!isByteBufferReadable(dis)) {
       // If InputStream does not support the ByteBuffer read, just read to heap and copy bytes to
       // the destination ByteBuff.
       byte[] heapBuf = new byte[length];
       IOUtils.readFully(dis, heapBuf, 0, length);
-      annotateHeapBytesRead(copy, length);
-      span.addEvent("BlockIOUtils.readFully", copy.build());
+      annotateHeapBytesRead(attributesBuilder, length);
+      span.addEvent("BlockIOUtils.readFully", attributesBuilder.build());
       copyToByteBuff(heapBuf, 0, length, buf);
       return;
     }
@@ -123,8 +125,8 @@ public static void readFully(ByteBuff buf, FSDataInputStream dis, int length,
         }
       }
     } finally {
-      annotateBytesRead(copy, directBytesRead, heapBytesRead);
-      span.addEvent("BlockIOUtils.readFully", copy.build());
+      annotateBytesRead(attributesBuilder, directBytesRead, heapBytesRead);
+      span.addEvent("BlockIOUtils.readFully", attributesBuilder.build());
     }
   }
 
@@ -136,13 +138,11 @@ public static void readFully(ByteBuff buf, FSDataInputStream dis, int length,
    * @param length to read
    * @throws IOException if any io error encountered.
    */
-  public static void readFullyWithHeapBuffer(InputStream in, ByteBuff out, int length,
-    AttributesBuilder attributesBuilder) throws IOException {
+  public static void readFullyWithHeapBuffer(InputStream in, ByteBuff out, int length)
+    throws IOException {
     if (length < 0) {
       throw new IllegalArgumentException("Length must not be negative: " + length);
     }
-    final Span span = Span.current();
-    final AttributesBuilder copy = Attributes.builder().putAll(attributesBuilder.build());
     int heapBytesRead = 0;
     int remain = length, count;
     byte[] buffer = new byte[1024];
@@ -158,8 +158,10 @@ public static void readFullyWithHeapBuffer(InputStream in, ByteBuff out, int len
         heapBytesRead += count;
       }
     } finally {
-      annotateHeapBytesRead(copy, heapBytesRead);
-      span.addEvent("BlockIOUtils.readFullyWithHeapBuffer", copy.build());
+      final Span span = Span.current();
+      final AttributesBuilder attributesBuilder = builderFromContext(Context.current());
+      annotateHeapBytesRead(attributesBuilder, heapBytesRead);
+      span.addEvent("BlockIOUtils.readFullyWithHeapBuffer", attributesBuilder.build());
     }
   }
 
@@ -177,9 +179,7 @@ public static void readFullyWithHeapBuffer(InputStream in, ByteBuff out, int len
    * @throws IOException if failed to read the necessary bytes
    */
   private static boolean readWithExtraOnHeap(InputStream in, byte[] buf, int bufOffset,
-    int necessaryLen, int extraLen, AttributesBuilder attributesBuilder) throws IOException {
-    final Span span = Span.current();
-    final AttributesBuilder copy = Attributes.builder().putAll(attributesBuilder.build());
+    int necessaryLen, int extraLen) throws IOException {
     int heapBytesRead = 0;
     int bytesRemaining = necessaryLen + extraLen;
     try {
@@ -199,8 +199,10 @@ private static boolean readWithExtraOnHeap(InputStream in, byte[] buf, int bufOf
         heapBytesRead += ret;
       }
     } finally {
-      annotateHeapBytesRead(copy, heapBytesRead);
-      span.addEvent("BlockIOUtils.readWithExtra", copy.build());
+      final Span span = Span.current();
+      final AttributesBuilder attributesBuilder = builderFromContext(Context.current());
+      annotateHeapBytesRead(attributesBuilder, heapBytesRead);
+      span.addEvent("BlockIOUtils.readWithExtra", attributesBuilder.build());
     }
     return bytesRemaining <= 0;
   }
@@ -217,17 +219,15 @@ private static boolean readWithExtraOnHeap(InputStream in, byte[] buf, int bufOf
    * @throws IOException if failed to read the necessary bytes.
    */
   public static boolean readWithExtra(ByteBuff buf, FSDataInputStream dis, int necessaryLen,
-    int extraLen, AttributesBuilder attributesBuilder) throws IOException {
+    int extraLen) throws IOException {
     if (!isByteBufferReadable(dis)) {
       // If InputStream does not support the ByteBuffer read, just read to heap and copy bytes to
       // the destination ByteBuff.
       byte[] heapBuf = new byte[necessaryLen + extraLen];
-      boolean ret = readWithExtraOnHeap(dis, heapBuf, 0, necessaryLen, extraLen, attributesBuilder);
+      boolean ret = readWithExtraOnHeap(dis, heapBuf, 0, necessaryLen, extraLen);
       copyToByteBuff(heapBuf, 0, heapBuf.length, buf);
       return ret;
     }
-    final Span span = Span.current();
-    final AttributesBuilder copy = Attributes.builder().putAll(attributesBuilder.build());
     int directBytesRead = 0, heapBytesRead = 0;
     ByteBuffer[] buffers = buf.nioByteBuffers();
     int bytesRead = 0;
@@ -259,8 +259,10 @@ public static boolean readWithExtra(ByteBuff buf, FSDataInputStream dis, int nec
         }
       }
     } finally {
-      annotateBytesRead(copy, directBytesRead, heapBytesRead);
-      span.addEvent("BlockIOUtils.readWithExtra", copy.build());
+      final Span span = Span.current();
+      final AttributesBuilder attributesBuilder = builderFromContext(Context.current());
+      annotateBytesRead(attributesBuilder, directBytesRead, heapBytesRead);
+      span.addEvent("BlockIOUtils.readWithExtra", attributesBuilder.build());
     }
     return (extraLen > 0) && (bytesRead == necessaryLen + extraLen);
   }
@@ -281,8 +283,8 @@ public static boolean readWithExtra(ByteBuff buf, FSDataInputStream dis, int nec
    * @throws IOException if failed to read the necessary bytes
    */
   public static boolean preadWithExtra(ByteBuff buff, FSDataInputStream dis, long position,
-    int necessaryLen, int extraLen, AttributesBuilder attributesBuilder) throws IOException {
-    return preadWithExtra(buff, dis, position, necessaryLen, extraLen, false, attributesBuilder);
+    int necessaryLen, int extraLen) throws IOException {
+    return preadWithExtra(buff, dis, position, necessaryLen, extraLen, false);
   }
 
   /**
@@ -302,24 +304,18 @@ public static boolean preadWithExtra(ByteBuff buff, FSDataInputStream dis, long
    * @throws IOException if failed to read the necessary bytes
    */
   public static boolean preadWithExtra(ByteBuff buff, FSDataInputStream dis, long position,
-    int necessaryLen, int extraLen, boolean readAllBytes, AttributesBuilder attributesBuilder)
-    throws IOException {
+    int necessaryLen, int extraLen, boolean readAllBytes) throws IOException {
     boolean preadbytebuffer = dis.hasCapability("in:preadbytebuffer");
 
     if (preadbytebuffer) {
-      return preadWithExtraDirectly(buff, dis, position, necessaryLen, extraLen, readAllBytes,
-        attributesBuilder);
+      return preadWithExtraDirectly(buff, dis, position, necessaryLen, extraLen, readAllBytes);
     } else {
-      return preadWithExtraOnHeap(buff, dis, position, necessaryLen, extraLen, readAllBytes,
-        attributesBuilder);
+      return preadWithExtraOnHeap(buff, dis, position, necessaryLen, extraLen, readAllBytes);
     }
   }
 
   private static boolean preadWithExtraOnHeap(ByteBuff buff, FSDataInputStream dis, long position,
-    int necessaryLen, int extraLen, boolean readAllBytes, AttributesBuilder attributesBuilder)
-    throws IOException {
-    final Span span = Span.current();
-    final AttributesBuilder copy = Attributes.builder().putAll(attributesBuilder.build());
+    int necessaryLen, int extraLen, boolean readAllBytes) throws IOException {
     int remain = necessaryLen + extraLen;
     byte[] buf = new byte[remain];
     int bytesRead = 0;
@@ -336,18 +332,17 @@ private static boolean preadWithExtraOnHeap(ByteBuff buff, FSDataInputStream dis
         remain -= ret;
       }
     } finally {
-      annotateHeapBytesRead(copy, bytesRead);
-      span.addEvent("BlockIOUtils.preadWithExtra", copy.build());
+      final Span span = Span.current();
+      final AttributesBuilder attributesBuilder = builderFromContext(Context.current());
+      annotateHeapBytesRead(attributesBuilder, bytesRead);
+      span.addEvent("BlockIOUtils.preadWithExtra", attributesBuilder.build());
     }
     copyToByteBuff(buf, 0, bytesRead, buff);
     return (extraLen > 0) && (bytesRead == necessaryLen + extraLen);
   }
 
   private static boolean preadWithExtraDirectly(ByteBuff buff, FSDataInputStream dis, long position,
-    int necessaryLen, int extraLen, boolean readAllBytes, AttributesBuilder attributesBuilder)
-    throws IOException {
-    final Span span = Span.current();
-    final AttributesBuilder copy = Attributes.builder().putAll(attributesBuilder.build());
+    int necessaryLen, int extraLen, boolean readAllBytes) throws IOException {
     int directBytesRead = 0, heapBytesRead = 0;
     int remain = necessaryLen + extraLen, bytesRead = 0, idx = 0;
     ByteBuffer[] buffers = buff.nioByteBuffers();
@@ -391,8 +386,10 @@ private static boolean preadWithExtraDirectly(ByteBuff buff, FSDataInputStream d
         }
       }
     } finally {
-      annotateBytesRead(copy, directBytesRead, heapBytesRead);
-      span.addEvent("BlockIOUtils.preadWithExtra", copy.build());
+      final Span span = Span.current();
+      final AttributesBuilder attributesBuilder = builderFromContext(Context.current());
+      annotateBytesRead(attributesBuilder, directBytesRead, heapBytesRead);
+      span.addEvent("BlockIOUtils.preadWithExtra", attributesBuilder.build());
     }
 
     return (extraLen > 0) && (bytesRead == necessaryLen + extraLen);
@@ -421,6 +418,18 @@ private static int copyToByteBuff(byte[] buf, int offset, int len, ByteBuff out)
     return len;
   }
 
+  /**
+   * Construct a fresh {@link AttributesBuilder} from the provided {@link Context}, populated with
+   * relevant attributes populated by {@link HFileContextAttributesBuilderConsumer#CONTEXT_KEY}.
+   */
+  private static AttributesBuilder builderFromContext(Context context) {
+    final AttributesBuilder attributesBuilder = Attributes.builder();
+    Optional.ofNullable(context)
+      .map(val -> val.get(HFileContextAttributesBuilderConsumer.CONTEXT_KEY))
+      .ifPresent(c -> c.accept(attributesBuilder));
+    return attributesBuilder;
+  }
+
   /**
    * Conditionally annotate {@code span} with the appropriate attribute when value is non-zero.
    */

hbase-common/src/main/java/org/apache/hadoop/hbase/trace/HBaseSemanticAttributes.java

@@ -112,7 +112,7 @@ public final class HBaseSemanticAttributes {
    * filesystem data.
    */
   public enum ReadType {
-    // TODO: promote this to the FSReader#readBlockData API.
+    // TODO: promote this to the FSReader#readBlockData API. Or somehow instead use Scan.ReadType.
     POSITIONAL_READ,
     SEEK_PLUS_READ,
   }