More powerful client service transformation

[mleclercq]

(edited to not re-add the R on client services)

I tried to add distributed tracing on the client side of a gRPC service and I think the available transformations are not powerful enough to cover my use case.

With the UIO[SafeMetadata], I can inject the tracing context to pass it to the server side, but I also would like to create a trace span on the client side with the some attributes from the gRPC call (see here for the recommendation from open-telemetry). To do so I need to apply a transformation on the effect of the call where I have access to the MethodDescriptor but with a Transform I do not have access to the MethodDescriptor (or CallOptions).

I've sketched up a possible solution in https://github.com/mleclercq/zio-grpc/blob/client-aspect/core/src/main/scala/scalapb/zio_grpc/clientAspects.scala

The idea is that the transformations (that I called ClientAspect) are applied on objects that wrap the actual call:

final case class CallEffect[Req, Res](
    call: (MethodDescriptor[Req, Res], CallOptions, UIO[SafeMetadata]) => ZIO[Any, StatusException, Res]
)

final case class CallStream[Req, Res](
    call: (MethodDescriptor[Req, Res], CallOptions, UIO[SafeMetadata]) => ZStream[Any, StatusException, Res]
)

trait ClientAspect { self =>
  // looks like a Transform, but can modify the full call, not just the effect that produces the result
  def effect[Req, Res](call: CallEffect[Req, Res]): CallEffect[Req, Res]
  def stream[Req, Res](call: CallStream[Req, Res]): CallStream[Req, Res]
}
object ClientAspect {
  // define standard aspects like setting deadlines etc.
}

Here is a mockup of what the generated code could look like:

class MyService(channel: ZChannel, aspects: ClientAspect = ClientAspect.identity)
    extends AspectTransformableClient[MyService] {

  val helloWorldMethod: MethodDescriptor[Int, String] = ???

  def helloWorld(req: Int): ZIO[Any, StatusException, String] = {
    // The "raw" CallEffect that does the actual call
    val callEffect = CallEffect[Int, String]((method, options, metadata) =>
      metadata.flatMap(
        scalapb.zio_grpc.client.ClientCalls.unaryCall(channel, method, options, _, req)
      )
    )

    // Modify the CallEffect with the aspects and call it
    aspects.effect(callEffect).call(helloWorldMethod, CallOptions.DEFAULT, SafeMetadata.make)

  }
  def transform(aspect: ClientAspect): MyService =
    new MyService(channel, aspects >>> aspect) {}
}

object MyService {
  def scoped(channel: ZManagedChannel, aspect: ClientAspect = ClientAspect.identity): ZIO[Scope, Throwable, MyService] =
    channel.map(new MyService(_, aspect))
}

The AspectTransformableClient trait extended by the generated MyService class is the equivalent of the existing ClientMethods and TransformableClient:

trait AspectTransformableClient[Repr] {
  def transform(aspect: ClientAspect): Repr

  def @@(aspect: ClientAspect): Repr = transform(aspect)

  def mapCallOptions(f: CallOptions => CallOptions): Repr        = transform(ClientAspect.mapCallOptions(f))
  def mapMetadataZIO(f: SafeMetadata => UIO[SafeMetadata]): Repr = transform(ClientAspect.mapMetadataZIO(f))
  def withMetadataZIO(metadata: UIO[SafeMetadata]): Repr         = transform(ClientAspect.withMetadataZIO(metadata))
  def withCallOptions(callOptions: CallOptions): Repr            = transform(ClientAspect.withCallOptions(callOptions))
  def withDeadline(deadline: Deadline): Repr                     = transform(ClientAspect.withDeadline(deadline))
  def withTimeout(duration: Duration): Repr                      = transform(ClientAspect.withTimeout(duration))
  def withTimeoutMillis(millis: Long): Repr                      = transform(ClientAspect.withTimeoutMillis(millis))
}

In the code linked above, I also define a mockup of a Tracer service and show how a GrpcClientTracer aspect could be defined.

This proposal significantly changes how the client service classes are generated (different constructor parameters) which I guess would be a problem for backward compatibility but I think it is be possible to generate new client services classes along side the existing ones.

Please let me know if you think this proposal makes sens and I could try to come up with an actual PR.

[thesamet]

Hi @mleclercq thanks for the suggestion. It is interesting. However, have you tried to implement the support for distributed tracing through ZClientInterceptor?

[mleclercq]

While I think it may be possible to do it with ZClientInterceptor, I didn't actually try because it doesn't provide an easy access to the call effect. I only have access to the ZClientCall which is quite low level so it would be difficult to make a trace span that covers the execution of the call (would need to plug into the listener to know when the call completed).

On the other hand, while implementing the server-side for distributed tracing, I realized that my ClientAspect is really a ZTransform in disguise 😄. On the server-side, the distributed tracing implements nicely as a ZTransform[RequestContext, RequestContext] and that made me think that I may be reinventing the wheel...

More precisely, a ClientAspect is just a ZTransform[CallContext, CallContext] where CallContext is defined as:

final case class CallContext(
    method: MethodDescriptor[_, _],
    options: CallOptions,
    metadata: UIO[SafeMetadata]
)

So we could modify TransformableClient as follow:

trait TransformableClient[Repr] {
  def transform(t: Transform): Repr = transform(t.toZTransform[CallContext])

  def transform(t: ZTransform[CallContext, CallContext]): Repr
}

Then ClientMethods.mapCallOptions and ClientMethods.mapMetadataZIO can be implemented in term of a ZTransform[CallContext, CallContext]:

def mapCallOptions(f: CallOptions => CallOptions): Repr =
  transform(ZTransform[CallContext, CallContext] { context =>
    ZIO.succeed(context.copy(options = f(context.options)))
  })

def mapMetadataZIO(f: SafeMetadata => UIO[SafeMetadata]): Repr =
  transform(ZTransform[CallContext, CallContext] { context =>
    ZIO.succeed(context.copy(metadata = context.metadata.flatMap(f)))
  })

And I think that code-gen could be simplified. The transform methods could be left abstract in the traits and only implemented in the concrete ServiceStub. mapCallOptions and mapMetadataZIO do not need to be generated since we have a generic implementation.

[thesamet]

Addressed by #528