refactor: make request context a little nicer to use (#12)

Author: aaron-steinfeld

grpc-client-utils/src/main/java/org/hypertrace/core/grpcutils/client/GrpcClientRequestContextUtil.java

@@ -1,6 +1,5 @@
 package org.hypertrace.core.grpcutils.client;
 
-import io.grpc.Context;
 import java.util.Map;
 import java.util.concurrent.Callable;
 import javax.annotation.Nonnull;
@@ -40,15 +39,7 @@ public static void executeInTenantContext(String tenantId, Runnable r) {
   public static <V> V executeWithHeadersContext(@Nonnull Map<String, String> headers, @Nonnull Callable<V> c) {
     RequestContext requestContext = new RequestContext();
     headers.forEach(requestContext::add);
-
-    try {
-      return Context.current().withValue(RequestContext.CURRENT, requestContext).call(c);
-    } catch (Exception e) {
-      if (e instanceof RuntimeException) {
-        throw (RuntimeException)e;
-      }
-      throw new RuntimeException(e);
-    }
+    return requestContext.call(c);
   }
 
   /**
@@ -60,7 +51,6 @@ public static <V> V executeWithHeadersContext(@Nonnull Map<String, String> heade
   public static void executeWithHeadersContext(@Nonnull Map<String, String> headers, @Nonnull Runnable r) {
     RequestContext requestContext = new RequestContext();
     headers.forEach(requestContext::add);
-
-    Context.current().withValue(RequestContext.CURRENT, requestContext).run(r);
+    requestContext.run(r);
   }
 }

grpc-context-utils/src/main/java/org/hypertrace/core/grpcutils/context/RequestContext.java

@@ -4,6 +4,8 @@
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Optional;
+import java.util.concurrent.Callable;
+import javax.annotation.Nonnull;
 
 /**
  * Context of the GRPC request that should be carried and can made available to the services so that
@@ -41,4 +43,19 @@ public Optional<String> get(String headerKey) {
   public Map<String, String> getAll() {
     return Map.copyOf(headers);
   }
+
+  public <V> V call(@Nonnull Callable<V> callable) {
+    try {
+      return Context.current().withValue(RequestContext.CURRENT, this).call(callable);
+    } catch (Exception e) {
+      if (e instanceof RuntimeException) {
+        throw (RuntimeException) e;
+      }
+      throw new RuntimeException(e);
+    }
+  }
+
+  public void run(@Nonnull Runnable runnable) {
+    Context.current().withValue(RequestContext.CURRENT, this).run(runnable);
+  }
 }