api: Add Int,Long, and String Tag variants (#242)

6 files changed, 185 insertions, 23 deletions

diff --git a/api/src/main/java/io/perfmark/Impl.java b/api/src/main/java/io/perfmark/Impl.java
index 7e33b6f..10a37d8 100644
--- a/api/src/main/java/io/perfmark/Impl.java
+++ b/api/src/main/java/io/perfmark/Impl.java
@@ -16,6 +16,10 @@
 
 package io.perfmark;
 
+import java.util.function.Function;
+import java.util.function.ToIntFunction;
+import java.util.function.ToLongFunction;
+
 public class Impl {
   static final String NO_TAG_NAME = "";
   static final long NO_TAG_ID = Long.MIN_VALUE;
@@ -42,7 +46,7 @@ public class Impl {
     return false;
   }
 
-  protected <T> void startTask(T taskNameObject, StringFunction<? super T> taskNameFunc) {}
+  protected <T> void startTask(T taskNameObject, Function<? super T, String> taskNameFunc) {}
 
   protected void startTask(String taskName, Tag tag) {}
 
@@ -79,7 +83,12 @@ public class Impl {
   protected void attachTag(String tagName, long tagValue0, long tagValue1) {}
 
   protected <T> void attachTag(
-      String tagName, T tagObject, StringFunction<? super T> stringFunction) {}
+      String tagName, T tagObject, Function<? super T, ? extends String> stringFunction) {}
+
+  protected <T> void attachTag(String tagName, T tagObject, ToIntFunction<? super T> intFunction) {}
+
+  protected <T> void attachTag(
+      String tagName, T tagObject, ToLongFunction<? super T> longFunction) {}
 
   protected Tag createTag(String tagName, long tagId) {
     return NO_TAG;
diff --git a/api/src/main/java/io/perfmark/PerfMark.java b/api/src/main/java/io/perfmark/PerfMark.java
index a064e98..6382feb 100644
--- a/api/src/main/java/io/perfmark/PerfMark.java
+++ b/api/src/main/java/io/perfmark/PerfMark.java
@@ -20,6 +20,9 @@ import com.google.errorprone.annotations.CheckReturnValue;
 import com.google.errorprone.annotations.DoNotCall;
 import com.google.errorprone.annotations.MustBeClosed;
 import java.lang.reflect.Method;
+import java.util.function.Function;
+import java.util.function.ToIntFunction;
+import java.util.function.ToLongFunction;
 
 /**
  * PerfMark is a very low overhead tracing library. To use PerfMark, annotate the code that needs to
@@ -546,7 +549,27 @@ public final class PerfMark {
    *
    * @param tagName The name of the value being attached
    * @param tagObject The tag object which will passed to the stringFunction.
-   * @param stringFunction The function that will convert the object to
+   * @param stringFunction The function that will convert the object to a tag
+   * @param <T> the type of tag object to be stringified
+   * @since 0.27.0
+   */
+  public static <T> void attachStringTag(
+      String tagName, T tagObject, Function<? super T, ? extends String> stringFunction) {
+    impl.attachTag(tagName, tagObject, stringFunction);
+  }
+
+  /**
+   * Attaches an additional keyed tag to the current active task. The tag provided is independent of
+   * the tag used with {@code startTask} and {@code stopTask}. This tag operation is different than
+   * {@link Tag} in that the tag value has an associated name (also called a key). The tag name and
+   * value are attached to the most recently started task, and don't have to match any other tags.
+   * This method is useful for when you have the tag information after the task is started.
+   *
+   * <p>Prefer {@link #attachStringTag(String, Object, Function)} over this one.
+   *
+   * @param tagName The name of the value being attached
+   * @param tagObject The tag object which will passed to the stringFunction.
+   * @param stringFunction The function that will convert the object to a tag
    * @param <T> the type of tag object to be stringified
    * @since 0.22.0
    */
@@ -555,6 +578,42 @@ public final class PerfMark {
     impl.attachTag(tagName, tagObject, stringFunction);
   }
 
+  /**
+   * Attaches an additional keyed tag to the current active task. The tag provided is independent of
+   * the tag used with {@code startTask} and {@code stopTask}. This tag operation is different than
+   * {@link Tag} in that the tag value has an associated name (also called a key). The tag name and
+   * value are attached to the most recently started task, and don't have to match any other tags.
+   * This method is useful for when you have the tag information after the task is started.
+   *
+   * @param tagName The name of the value being attached
+   * @param tagObject The tag object which will passed to the intFunction.
+   * @param intFunction The function that will convert the object to a tag
+   * @param <T> the type of tag object to mapped to in.
+   * @since 0.27.0
+   */
+  public static <T> void attachIntTag(
+      String tagName, T tagObject, ToIntFunction<? super T> intFunction) {
+    impl.attachTag(tagName, tagObject, intFunction);
+  }
+
+  /**
+   * Attaches an additional keyed tag to the current active task. The tag provided is independent of
+   * the tag used with {@code startTask} and {@code stopTask}. This tag operation is different than
+   * {@link Tag} in that the tag value has an associated name (also called a key). The tag name and
+   * value are attached to the most recently started task, and don't have to match any other tags.
+   * This method is useful for when you have the tag information after the task is started.
+   *
+   * @param tagName The name of the value being attached
+   * @param tagObject The tag object which will passed to the intFunction.
+   * @param longFunction The function that will convert the object to a tag
+   * @param <T> the type of tag object to mapped to in.
+   * @since 0.27.0
+   */
+  public static <T> void attachLongTag(
+      String tagName, T tagObject, ToLongFunction<? super T> longFunction) {
+    impl.attachTag(tagName, tagObject, longFunction);
+  }
+
   private static final Impl impl;
 
   static {
diff --git a/api/src/main/java/io/perfmark/StringFunction.java b/api/src/main/java/io/perfmark/StringFunction.java
index 19506e0..8360af0 100644
--- a/api/src/main/java/io/perfmark/StringFunction.java
+++ b/api/src/main/java/io/perfmark/StringFunction.java
@@ -16,6 +16,8 @@
 
 package io.perfmark;
 
+import java.util.function.Function;
+
 /**
  * This interface is equivalent to {@code java.util.function.Function}. It is here as a
  * compatibility shim to make PerfMark compatible with Java 6. This will likely be removed if
@@ -24,7 +26,8 @@ package io.perfmark;
  * @since 0.22.0
  * @param <T> The type to turn into a String.
  */
-public interface StringFunction<T> {
+@FunctionalInterface
+public interface StringFunction<T> extends Function<T, String> {
 
   /**
    * Takes the given argument and produces a String.
@@ -33,5 +36,6 @@ public interface StringFunction<T> {
    * @param t the subject to Stringify
    * @return the String
    */
+  @Override
   String apply(T t);
 }
diff --git a/api/src/test/java/io/perfmark/CompatibilityTest.java b/api/src/test/java/io/perfmark/CompatibilityTest.java
index 3a6013d..3381bd9 100644
--- a/api/src/test/java/io/perfmark/CompatibilityTest.java
+++ b/api/src/test/java/io/perfmark/CompatibilityTest.java
@@ -420,6 +420,26 @@ public class CompatibilityTest {
     assertThat(marks).hasSize(3);
   }
 
+  @Test
+  public void implOverride() throws Exception {
+    Assume.assumeTrue(minorVersion >= 17);
+
+    Class<?> implClass = Class.forName("io.perfmark.Impl", false, perfMarkClz.getClassLoader());
+    Class<?> secretImplClass =
+        Class.forName(
+            "io.perfmark.impl.SecretPerfMarkImpl$PerfMarkImpl",
+            false,
+            perfMarkClz.getClassLoader());
+
+    for (Method method : implClass.getDeclaredMethods()) {
+      if ((method.getModifiers() & Modifier.STATIC) != 0) {
+        continue;
+      }
+      Method m2 = secretImplClass.getDeclaredMethod(method.getName(), method.getParameterTypes());
+      assertThat(m2).isNotNull();
+    }
+  }
+
   private final class ApiOverrideClassLoader extends ClassLoader {
     @Override
     protected Class<?> loadClass(String name, boolean resolve) throws ClassNotFoundException {
diff --git a/api/src/test/java/io/perfmark/PerfMarkTest.java b/api/src/test/java/io/perfmark/PerfMarkTest.java
index 645edf3..0eac672 100644
--- a/api/src/test/java/io/perfmark/PerfMarkTest.java
+++ b/api/src/test/java/io/perfmark/PerfMarkTest.java
@@ -114,6 +114,9 @@ public class PerfMarkTest {
     PerfMark.startTask("task5", String::valueOf);
     PerfMark.attachTag(PerfMark.createTag("extra"));
     PerfMark.attachTag("name", "extra2", String::valueOf);
+    PerfMark.attachStringTag("name", "extra3", String::valueOf);
+    PerfMark.attachIntTag("name", List.of(), List::size);
+    PerfMark.attachLongTag("name", 2d, Double::longValue);
     Link link = PerfMark.linkOut();
     PerfMark.linkIn(link);
     PerfMark.stopTask();
@@ -127,7 +130,7 @@ public class PerfMarkTest {
 
     List<Mark> marks = new ArrayList<>(Storage.readForTest());
 
-    Truth.assertThat(marks).hasSize(24);
+    Truth.assertThat(marks).hasSize(27);
     List<Mark> expected =
         Arrays.asList(
             Mark.taskStart(gen, marks.get(0).getNanoTime(), "task1"),
@@ -140,20 +143,23 @@ public class PerfMarkTest {
             Mark.taskStart(gen, marks.get(7).getNanoTime(), "task5"),
             Mark.tag(gen, "extra", NO_TAG_ID),
             Mark.keyedTag(gen, "name", "extra2"),
+            Mark.keyedTag(gen, "name", "extra3"),
+            Mark.keyedTag(gen, "name", 0),
+            Mark.keyedTag(gen, "name", 2),
             Mark.link(gen, link.linkId),
             Mark.link(gen, -link.linkId),
-            Mark.taskEnd(gen, marks.get(12).getNanoTime()),
-            Mark.taskEnd(gen, marks.get(13).getNanoTime(), "task4"),
+            Mark.taskEnd(gen, marks.get(15).getNanoTime()),
+            Mark.taskEnd(gen, marks.get(16).getNanoTime(), "task4"),
             Mark.tag(gen, tag3.tagName, tag3.tagId),
-            Mark.taskEnd(gen, marks.get(15).getNanoTime(), "task3"),
+            Mark.taskEnd(gen, marks.get(18).getNanoTime(), "task3"),
             Mark.tag(gen, tag2.tagName, tag2.tagId),
-            Mark.taskEnd(gen, marks.get(17).getNanoTime(), "task2"),
+            Mark.taskEnd(gen, marks.get(20).getNanoTime(), "task2"),
             Mark.tag(gen, tag1.tagName, tag1.tagId),
-            Mark.taskEnd(gen, marks.get(19).getNanoTime(), "task1"),
-            Mark.taskStart(gen, marks.get(20).getNanoTime(), "task6"),
-            Mark.taskStart(gen, marks.get(21).getNanoTime(), "task7"),
-            Mark.taskEnd(gen, marks.get(22).getNanoTime()),
-            Mark.taskEnd(gen, marks.get(23).getNanoTime()));
+            Mark.taskEnd(gen, marks.get(22).getNanoTime(), "task1"),
+            Mark.taskStart(gen, marks.get(23).getNanoTime(), "task6"),
+            Mark.taskStart(gen, marks.get(24).getNanoTime(), "task7"),
+            Mark.taskEnd(gen, marks.get(25).getNanoTime()),
+            Mark.taskEnd(gen, marks.get(26).getNanoTime()));
     assertEquals(expected, marks);
   }
 
@@ -162,7 +168,8 @@ public class PerfMarkTest {
     Storage.resetForThread();
     PerfMark.setEnabled(true);
 
-    PerfMark.attachTag("name", "extra2", null);
+    StringFunction<String> nullFunction = null;
+    PerfMark.attachTag("name", "extra2", nullFunction);
 
     List<Mark> marks = Storage.readForTest();
     Truth.assertThat(marks).hasSize(1);
diff --git a/impl/src/main/java/io/perfmark/impl/SecretPerfMarkImpl.java b/impl/src/main/java/io/perfmark/impl/SecretPerfMarkImpl.java
index 173572d..4967b1a 100644
--- a/impl/src/main/java/io/perfmark/impl/SecretPerfMarkImpl.java
+++ b/impl/src/main/java/io/perfmark/impl/SecretPerfMarkImpl.java
@@ -21,6 +21,9 @@ import io.perfmark.Link;
 import io.perfmark.StringFunction;
 import io.perfmark.Tag;
 import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Function;
+import java.util.function.ToIntFunction;
+import java.util.function.ToLongFunction;
 import java.util.logging.Level;
 import java.util.logging.LogRecord;
 import java.util.logging.Logger;
@@ -251,8 +254,16 @@ final class SecretPerfMarkImpl {
       markRecorder.start(gen, taskName, subTaskName);
     }
 
-    @Override
+    /**
+     * This method is needed to work with old version of perfmark-api.
+     */
     protected <T> void startTask(T taskNameObject, StringFunction<? super T> stringFunction) {
+      Function<? super T, String> function = stringFunction;
+      startTask(taskNameObject, function);
+    }
+
+    @Override
+    protected <T> void startTask(T taskNameObject, Function<? super T, String> stringFunction) {
       final long gen = getGen();
       if (!isEnabled(gen)) {
         return;
@@ -349,9 +360,18 @@ final class SecretPerfMarkImpl {
       markRecorder.attachKeyedTag(gen, tagName, tagValue);
     }
 
+    /**
+     * This method is needed to work with old version of perfmark-api.
+     */
+    public <T> void attachTag(
+        String tagName, T tagObject, StringFunction<? super T> stringFunction) {
+      Function<? super T, ? extends String> function = stringFunction;
+      attachTag(tagName, tagObject, function);
+    }
+
     @Override
     protected <T> void attachTag(
-        String tagName, T tagObject, StringFunction<? super T> stringFunction) {
+        String tagName, T tagObject, Function<? super T, ? extends String> stringFunction) {
       final long gen = getGen();
       if (!isEnabled(gen)) {
         return;
@@ -360,8 +380,30 @@ final class SecretPerfMarkImpl {
       markRecorder.attachKeyedTag(gen, tagName, tagValue);
     }
 
+    @Override
+    protected <T> void attachTag(
+        String tagName, T tagObject, ToIntFunction<? super T> intFunction) {
+      final long gen = getGen();
+      if (!isEnabled(gen)) {
+        return;
+      }
+      long tagValue = deriveTagValue(tagName, tagObject, intFunction);
+      markRecorder.attachKeyedTag(gen, tagName, tagValue);
+    }
+
+    @Override
+    protected <T> void attachTag(
+        String tagName, T tagObject, ToLongFunction<? super T> longFunction) {
+      final long gen = getGen();
+      if (!isEnabled(gen)) {
+        return;
+      }
+      long tagValue = deriveTagValue(tagName, tagObject, longFunction);
+      markRecorder.attachKeyedTag(gen, tagName, tagValue);
+    }
+
     static <T> String deriveTagValue(
-        String tagName, T tagObject, StringFunction<? super T> stringFunction) {
+        String tagName, T tagObject, Function<? super T, ? extends String> stringFunction) {
       try {
         return stringFunction.apply(tagObject);
       } catch (Throwable t) {
@@ -370,7 +412,28 @@ final class SecretPerfMarkImpl {
       }
     }
 
-    static <T> String deriveTaskValue(T taskNameObject, StringFunction<? super T> stringFunction) {
+    static <T> long deriveTagValue(
+        String tagName, T tagNameObject, ToIntFunction<? super T> intFunction) {
+      try {
+        // implicit cast
+        return intFunction.applyAsInt(tagNameObject);
+      } catch (Throwable t) {
+        handleTagValueFailure(tagName, tagNameObject, intFunction, t);
+        return Mark.NO_TAG_ID;
+      }
+    }
+
+    static <T> long deriveTagValue(
+        String tagName, T tagNameObject, ToLongFunction<? super T> longFunction) {
+      try {
+        return longFunction.applyAsLong(tagNameObject);
+      } catch (Throwable t) {
+        handleTagValueFailure(tagName, tagNameObject, longFunction, t);
+        return Mark.NO_TAG_ID;
+      }
+    }
+
+    static <T> String deriveTaskValue(T taskNameObject, Function<? super T, String> stringFunction) {
       try {
         return stringFunction.apply(taskNameObject);
       } catch (Throwable t) {
@@ -380,7 +443,7 @@ final class SecretPerfMarkImpl {
     }
 
     static <T> void handleTagValueFailure(
-        String tagName, T tagObject, StringFunction<? super T> stringFunction, Throwable cause) {
+        String tagName, T tagObject, Object stringFunction, Throwable cause) {
       if (logger == null) {
         return;
       }
@@ -407,7 +470,7 @@ final class SecretPerfMarkImpl {
     }
 
     static <T> void handleTaskNameFailure(
-        T taskNameObject, StringFunction<? super T> stringFunction, Throwable cause) {
+        T taskNameObject, Object function, Throwable cause) {
       if (logger == null) {
         return;
       }
@@ -416,8 +479,8 @@ final class SecretPerfMarkImpl {
         if (localLogger.isLoggable(Level.FINE)) {
           LogRecord lr =
               new LogRecord(
-                  Level.FINE, "PerfMark.startTask failed: taskObject={0}, stringFunction={1}");
-          lr.setParameters(new Object[] {taskNameObject, stringFunction});
+                  Level.FINE, "PerfMark.startTask failed: taskObject={0}, function={1}");
+          lr.setParameters(new Object[] {taskNameObject, function});
           lr.setThrown(cause);
           localLogger.log(lr);
         }