Hooks interface & otel integration

[slinkydeveloper]

Introduce new hooks interface that will allow us to provide better integration with tracing, and other telemetry libs:

type AttemptResult = {
  type: "success"
} | {
  type: "retryableError"
  error: Error
} | {
  type: "terminalError"
  error: TerminalError
}

type Hooks = {
  // Set this hook to wrap the handler execution. Calling handlerRunner will cause the handler to run. 
  // Use it to propagate an async context storage.
  runInContext?: <T>(handlerRunner: () => Promise<T>) => Promise<T>;
  
  // Called when the attempt completes. Errors thrown inside this function are not propagated back to restate
  onAttemptEnd?: (result: AttemptResult) => void;
};

// If the hook provider fails, same rules as handler failures apply:
// * if fails with terminal error, invocation is terminated with terminal error
// * for other failures, it gets retried
type HooksProvider = (ctx: {
  serviceName: string;
  handlerName: string;
  key?: string;
  invocationId: string;
  request: Request;
}) => Hooks;

One can set these HooksProviders on a handler, on a service, or for the whole endpoint, and hooks will compose.

Tracing hooks provider example:

// Example provider: otel tracing
function otelTracing(tracer: Tracer): HooksProvider {
  return (ctx) => {
    // .... extract parent trace etc ...
    const span = tracer.startSpan(`${ctx.serviceName}/${ctx.handlerName}`, {"restate.invocation.id": ctx.invocationId});

    return {
      runInContext: (handlerRunner) =>
          otelContext.with(trace.setSpan(otelContext.active(), span), handlerRunner),
      onAttemptEnd: (result) => {
        if (result.type === "retryableError") {
          span.recordException(result.error);
          span.setAttribute("transient", true);
        } else if (result.type === "terminalError") {
          span.recordException(result.error);
        }
        span.end();
      },
    };
  };
}

const greeter = restate.service({
  name: "Greeter",
  handlers: {
    greet: async (ctx: restate.Context, name: string) => {
      // otel context — via otel API
      const span = trace.getActiveSpan();
      span?.addEvent("greeting_started");

      return `Hello, ${name}!`;
    },
  },
  options: {
    hooks: [otelTracing(tracer)]
  }
});

We could extend these hooks in future to do more things, like hooks on outgoing requests, etc.

[srujun]

Overall looks pretty good, just a few follow-up questions about this:

1. Will the hook also include attemptHeaders within the Request object?
2. Can this also expose the extracted OTel Context directly in the hook provider input, instead of the user needing to implement propagation.extract()?
3. Will ctx.run() spans be created within the active OTel context on the handler side?

[slinkydeveloper]

Will the hook also include attemptHeaders within the Request object?

Yes

Can this also expose the extracted OTel Context directly in the hook provider input, instead of the user needing to implement propagation.extract()?

WDYM by hook provider input here? In the example, we propagate using the otel api directly.

Will ctx.run() spans be created within the active OTel context on the handler side?

That's perhaps something we can add to the various hooks, like:

type Hooks = {

  // -- Other stuff

  wrapRun?: (runName: string, runClosure: () => Promise<void>) => Promise<void>;
};

So one can create a span, starting from the active span, and propagate it inside when calling runClosure. Something like that.

[mupperton]

Do you think this hooks interface would be a good place to have a wrapCall and wrapSend in addition to wrapRun?

This feels like the right successor to my attempt at this in #539 - so I can customise the rpc.opts / rpc.sendOpts usage on all RPC calls/sends at a handler/service/deployment level using options.hooks

[slinkydeveloper]

I guess we can also add that as well @mupperton