Add ServiceDefinition to OperationContext

nexus-sdk/src/main/java/io/nexusrpc/handler/OperationContext.java

@@ -1,6 +1,7 @@
 package io.nexusrpc.handler;
 
 import io.nexusrpc.Link;
+import io.nexusrpc.ServiceDefinition;
 import java.time.Instant;
 import java.util.*;
 import java.util.stream.Collectors;
@@ -23,20 +24,25 @@ public static Builder newBuilder(OperationContext context) {
   private final Map<String, String> headers;
   // This is not included in equals, hashCode, or toString
   private final @Nullable OperationMethodCanceller methodCanceller;
-  private final List<Link> links = new ArrayList<>();
   private final Instant deadline;
+  private final List<Link> links;
+  private final @Nullable ServiceDefinition serviceDefinition;
 
   private OperationContext(
       String service,
       String operation,
       Map<String, String> headers,
       @Nullable OperationMethodCanceller methodCanceller,
-      Instant deadline) {
+      Instant deadline,
+      @Nullable ServiceDefinition serviceDefinition,
+      List<Link> links) {
     this.service = service;
     this.operation = operation;
     this.headers = headers;
     this.methodCanceller = methodCanceller;
     this.deadline = deadline;
+    this.serviceDefinition = serviceDefinition;
+    this.links = links;
   }
 
   /** Service name for the call. */
@@ -75,6 +81,11 @@ public boolean isMethodCancelled() {
     return methodCanceller == null ? null : methodCanceller.getCancellationReason();
   }
 
+  /** Get the service definition associated with this operation context, if any. */
+  public @Nullable ServiceDefinition getServiceDefinition() {
+    return serviceDefinition;
+  }
+
   /**
    * Get the deadline for the operation handler method. This is the time by which the method should
    * complete. This is not the operation's deadline.
@@ -136,18 +147,21 @@ public OperationContext removeMethodCancellationListener(
 
   @Override
   public boolean equals(Object o) {
-    if (this == o) return true;
     if (o == null || getClass() != o.getClass()) return false;
     OperationContext that = (OperationContext) o;
     return Objects.equals(service, that.service)
         && Objects.equals(operation, that.operation)
         && Objects.equals(headers, that.headers)
-        && Objects.equals(links, that.links);
+        && Objects.equals(methodCanceller, that.methodCanceller)
+        && Objects.equals(deadline, that.deadline)
+        && Objects.equals(links, that.links)
+        && Objects.equals(serviceDefinition, that.serviceDefinition);
   }
 
   @Override
   public int hashCode() {
-    return Objects.hash(service, operation, headers);
+    return Objects.hash(
+        service, operation, headers, methodCanceller, deadline, links, serviceDefinition);
   }
 
   @Override
@@ -161,8 +175,14 @@ public String toString() {
         + '\''
         + ", headers="
         + headers
+        + ", methodCanceller="
+        + methodCanceller
+        + ", deadline="
+        + deadline
         + ", links="
         + links
+        + ", serviceDefinition="
+        + serviceDefinition
         + '}';
   }
 
@@ -173,15 +193,25 @@ public static class Builder {
     private final SortedMap<String, String> headers;
     private @Nullable OperationMethodCanceller methodCanceller;
     private @Nullable Instant deadline;
+    private @Nullable ServiceDefinition serviceDefinition;
+    // Currently links are not set in the builder, but they need to be passed though to go from
+    // OperationContext -> Builder
+    // and back to OperationContext, so we keep them here.
+    private final List<Link> links;
 
     private Builder() {
       headers = new TreeMap<>(String.CASE_INSENSITIVE_ORDER);
+      links = new ArrayList<>();
     }
 
     private Builder(OperationContext context) {
       service = context.service;
       operation = context.operation;
       headers = new TreeMap<>(context.headers);
+      methodCanceller = context.methodCanceller;
+      deadline = context.deadline;
+      serviceDefinition = context.serviceDefinition;
+      links = context.links;
     }
 
     /** Set service. Required. */
@@ -219,6 +249,12 @@ public Builder setDeadline(Instant deadline) {
       return this;
     }
 
+    /** Sets the service definition. */
+    public Builder setServiceDefinition(ServiceDefinition serviceDefinition) {
+      this.serviceDefinition = serviceDefinition;
+      return this;
+    }
+
     /** Build the context. */
     public OperationContext build() {
       Objects.requireNonNull(service, "Service required");
@@ -236,7 +272,9 @@ public OperationContext build() {
           operation,
           Collections.unmodifiableMap(new TreeMap<>(normalizedHeaders)),
           methodCanceller,
-          deadline);
+          deadline,
+          serviceDefinition,
+          links);
     }
   }
 }

nexus-sdk/src/main/java/io/nexusrpc/handler/ServiceHandler.java

@@ -67,8 +67,12 @@ public OperationStartResult<HandlerResultContent> startOperation(
     if (handler == null) {
       throw newUnrecognizedOperationException(context.getService(), context.getOperation());
     }
+    // Populate the service definition in the context so that the handler can use it
+    OperationContext contextWithServiceDef =
+        OperationContext.newBuilder(context).setServiceDefinition(instance.getDefinition()).build();
+
     OperationHandler<Object, Object> interceptedHandler =
-        interceptOperationHandler(context, handler);
+        interceptOperationHandler(contextWithServiceDef, handler);
     OperationDefinition definition =
         instance.getDefinition().getOperations().get(context.getOperation());
 
@@ -84,7 +88,8 @@ public OperationStartResult<HandlerResultContent> startOperation(
     }
 
     // Invoke handler
-    OperationStartResult<?> result = interceptedHandler.start(context, details, inputObject);
+    OperationStartResult<?> result =
+        interceptedHandler.start(contextWithServiceDef, details, inputObject);
 
     // If the result is an async result we can just return, but if it's a sync result we need to
     // serialize back out to bytes
@@ -108,7 +113,13 @@ public HandlerResultContent fetchOperationResult(
     if (handler == null) {
       throw newUnrecognizedOperationException(context.getService(), context.getOperation());
     }
-    Object result = interceptOperationHandler(context, handler).fetchResult(context, details);
+    // Populate the service definition in the context so that the handler can use it
+    OperationContext contextWithServiceDef =
+        OperationContext.newBuilder(context).setServiceDefinition(instance.getDefinition()).build();
+
+    Object result =
+        interceptOperationHandler(contextWithServiceDef, handler)
+            .fetchResult(contextWithServiceDef, details);
     return resultToContent(result);
   }
 
@@ -136,7 +147,11 @@ public OperationInfo fetchOperationInfo(
     if (handler == null) {
       throw newUnrecognizedOperationException(context.getService(), context.getOperation());
     }
-    return interceptOperationHandler(context, handler).fetchInfo(context, details);
+    // Populate the service definition in the context so that the handler can use it
+    OperationContext contextWithServiceDef =
+        OperationContext.newBuilder(context).setServiceDefinition(instance.getDefinition()).build();
+    return interceptOperationHandler(contextWithServiceDef, handler)
+        .fetchInfo(contextWithServiceDef, details);
   }
 
   @Override
@@ -150,7 +165,11 @@ public void cancelOperation(OperationContext context, OperationCancelDetails det
     if (handler == null) {
       throw newUnrecognizedOperationException(context.getService(), context.getOperation());
     }
-    interceptOperationHandler(context, handler).cancel(context, details);
+    // Populate the service definition in the context so that the handler can use it
+    OperationContext contextWithServiceDef =
+        OperationContext.newBuilder(context).setServiceDefinition(instance.getDefinition()).build();
+    interceptOperationHandler(contextWithServiceDef, handler)
+        .cancel(contextWithServiceDef, details);
   }
 
   private static HandlerException newUnrecognizedOperationException(

nexus-sdk/src/test/java/io/nexusrpc/handler/OperationContextTest.java

@@ -3,13 +3,23 @@
 import static org.junit.jupiter.api.Assertions.assertEquals;
 
 import io.nexusrpc.Link;
+import io.nexusrpc.example.TestServices;
 import java.net.URI;
 import java.net.URISyntaxException;
 import java.time.Instant;
 import java.util.Arrays;
 import org.junit.jupiter.api.Test;
 
 public class OperationContextTest {
+
+  @ServiceImpl(service = TestServices.IntegerService.class)
+  public class IntegerServiceImpl {
+    @OperationImpl
+    public OperationHandler<Integer, Integer> operation() {
+      return OperationHandler.sync((ctx, details, input) -> 0);
+    }
+  }
+
   @Test
   void linkTest() throws URISyntaxException {
     OperationContext octx =
@@ -34,4 +44,23 @@ void deadlineTest() {
             .build();
     assertEquals(deadline, octx.getDeadline());
   }
+
+  @Test
+  void contextBuilderTest() throws URISyntaxException {
+    OperationMethodCanceller omc = new OperationMethodCanceller();
+    ServiceImplInstance serviceImpl = ServiceImplInstance.fromInstance(new IntegerServiceImpl());
+    Instant deadline = Instant.now().plusMillis(1000);
+    OperationContext octx =
+        OperationContext.newBuilder()
+            .setService("service")
+            .setOperation("operation")
+            .setDeadline(deadline)
+            .putHeader("key", "value")
+            .setMethodCanceller(omc)
+            .setServiceDefinition(serviceImpl.getDefinition())
+            .build();
+    URI url = new URI("http://somepath?k=v");
+    octx.setLinks(Link.newBuilder().setUri(url).setType("com.example.MyResource").build());
+    assertEquals(octx, OperationContext.newBuilder(octx).build());
+  }
 }