Add context propagation support #783

rgrandl · 2024-07-17T20:51:44Z

In the current implementation, weaver doesn't allow the user to propagate context information. We recommend users to define a struct that encapsulates the metadata information and add it as an argument to the method. However, more and more users are asking for an option to propagate metadata information using the context. This request comes especially from users that are using gRPC to communicate between their services, and gRPC provides a way to propagate metadata information using the context.

This PR implements a similar approach to gRPC [1] to propagate metadata information. We allow users to create a context where they can pass metadata as a map[string][]string.

E.g.,

...
ctx := context.Background()
meta := metadata.New(map[string][]string{"name":{"foo", "bar"}})
ctx = metadata.NewContext(ctx, meta)

dst.Get().Update(ctx, ....)

...
func( d *destination) Update(ctx, ...) (error) {
  meta := metadata.FromContext(ctx)
  vals := meta.Get("name")

  // Can call another component method with an updated context
  meta.Set("account", "x", "y")
  ctx = metadata.NewContext(ctx, meta)
  acct.Get().UpdateAccount(ctx, ...)
}

The user APIs are as follows:

// New creates a Metadata object that has operations
// to manipulate the metadata (e.g., Get, Set, Append, Delete.
func New(m map[string][]string) Metadata

// NewContext creates a new context that contains metadata meta, with
// parent context ctx.
func NewContext(ctx context.Context, meta Metadata) context.Context

// FromContext returns the metadata from the context. If the context
// doesn't have any metata, it returns an empty metadata.
func FromContext(ctx context.Context) Metadata

Note that similar to the gRPC metadata, we make sure the metadata keys are lowercased, because the user can create metadata w/o using our New() metadata constructor, hence they can set the keys however they want.

[1] https://pkg.go.dev/google.golang.org/grpc/metadata

rgrandl · 2024-07-17T20:52:51Z

@ghemawat do you mind taking a look? Tests are missing, but I would like to get your opinion on the APIs before I go ahead and polish the PR.

rgrandl

Thanks Sanjay. Updated the PR to incorporate your feedback.

ghemawat

API looks good.

Can you update the commit message.

ghemawat · 2024-07-18T14:33:05Z

weavertest/internal/simple/simple_test.go

+	// component. Because the UpdateMetadata and GetMetadata methods are not
+	// routed, we can end up updating the metadata at replica 0, and reading
+	// the metadata from replica 1. To avoid this scenario we call the
+	// UpdateMetadata method twice (one for each replica).


How are we sure that the two calls won't both go to the same replica?

Perhaps we should make it a bit better by calling it N times, where N is 10 for multi and 1 for others

Good point. Done

ghemawat · 2024-07-18T14:33:42Z

weavertest/internal/simple/simple_test.go

+			if err := dst.UpdateMetadata(ctx); err != nil {
+				t.Fatal(err)
+			}
+			if runner.Name == weavertest.Multi.Name {


The multi-replica comment belongs here, not above.

ghemawat · 2024-07-18T14:35:49Z

weavertest/internal/simple/simple_test.go

+			if err := dst.UpdateMetadata(ctx); err != nil {
+				t.Fatal(err)
+			}
+			if runner.Name == weavertest.Multi.Name {


Repeated code. Maybe add a helper: setAndGetMetadata(map[string]string) map[string]string

I changed the code a little bit. Does it look better?

ghemawat · 2024-07-18T14:37:21Z

godeps.txt

@@ -983,6 +984,9 @@ github.com/ServiceWeaver/weaver/sim/internal/bank
    go.opentelemetry.io/otel/codes
    go.opentelemetry.io/otel/trace
    reflect
+github.com/ServiceWeaver/weaver/weavermetadata


Suggested rename of package "callmetadata".

Caller will look something like:

ctx = callmetadata.NewContext(ctx, map[string]string{...})

Or perhaps just "metadata", which is what gRPC does.

Called it "metadata".

ghemawat · 2024-07-18T14:41:51Z

internal/net/call/call.go

-	msgHeaderSize = 16 + 8 + traceHeaderLen // handler_key + deadline + trace_context
-)
+// Size of the header included in each message.
+const msgHeaderSize = 16 + 8 + traceHeaderLen + metadataHeaderLen


Documentation of format should be updated above.

We don't use a fixed header anymore, so all these constants are gone.

ghemawat · 2024-07-18T14:54:24Z

internal/net/call/metadata.go

+
+// readContextMetadata returns the context metadata (if any).
+func readContextMetadata(meta []byte) map[string]string {
+	dec := codegen.NewDecoder(meta)


I wonder if we should be using codegen encoder/decoder to format the header completely. I.e., writeHeader that does something like:

put request id put payload length put trace boolean if present and then encode trace put #metadata entries put metadata entries

ghemawat · 2024-07-18T14:58:58Z

internal/net/call/call.go

-	// Size of the header included in each message.
-	msgHeaderSize = 16 + 8 + traceHeaderLen // handler_key + deadline + trace_context
-)
+// Size of the header included in each message.


Since this is a format change, are there any versions do we need to update? Do we use "call.go" for any of our internal communication where the two sides may be at different versions.

No, I think everything should work as expected

ghemawat · 2024-07-18T15:00:14Z

internal/net/call/trace.go

@@ -20,7 +20,7 @@ import (
 	"go.opentelemetry.io/otel/trace"
 )

-const traceHeaderLen = 25
+const traceHeaderLen = 25 // handler_key + deadline + trace_context


It would be nice if all of the header constants that depend on each other were described next to each other in the same constant block

ghemawat · 2024-07-18T15:03:29Z

weavermetadata/metadata.go

+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Package weavermetadata define the structure of the metadata supported by weaver.


Reword:

Package weavermetadata provides support for the propagation of metadata information from a component method caller to the callee. The metadata is propagated to the callee even if the caller and callee are not colocated in the same process.

The metadata is a map from string to string stored in context.Context. The map can be added to a context by calling NewContext.

... include example here ...

The metadata map can be extracted from a context by calling FromContext:

... include example here ...

ghemawat · 2024-07-18T15:04:40Z

weavermetadata/metadata.go

+// more complicated types, it's their responsibility to encode/decode these types
+// to/from string values.
+//
+// Note that all keys are automatically lowercase. This ensures the user avoids


We should drop this lower-casing I think.

rgrandl

Thanks, Sanjay. Very good suggestions.

ghemawat · 2024-08-06T15:45:04Z

internal/net/call/call.go

-	var hdr [msgHeaderSize]byte
-	copy(hdr[0:], h[:])
+	enc := codegen.NewEncoder()
+	enc.Bytes(h[:])


Bytes() seems suitable for variable-length data - it encodes the length. How about using Encoder.Grow and Decoder.Read to write exactly the correct number of bytes.

It also means that we won't have to worry about the wrong number of bytes showing up at the reader.

Good point. Changed.

ghemawat · 2024-08-06T15:45:55Z

internal/net/call/call.go

+	// Extract trace context and create a new child span to trace the method
+	// call on the server.
+	span := trace.SpanFromContext(ctx) // noop span
+	if sc := readTraceContext(dec); sc != nil && sc.IsValid() {


What if sc is invalid? Should we not fail the RPC with an error?

Good point. We are not failing the RPC right now if sc is invalid, but I think we should do that. Changed.

ghemawat · 2024-08-06T15:49:22Z

internal/net/call/call.go

+	// Encode the length of the header. Note that the header can have variable length,
+	// so we need to know at which offset in the encoded message starts the payload
+	// when we decode the message at the receiver.
+	enc.Int(len(enc.Data()))


enc.Uint? to avoid having to deal with negative number errors?

We don't need to encode the length of the header based on your suggestions.

ghemawat · 2024-08-06T15:51:31Z

internal/net/call/call.go

+	ctx = readContextMetadata(ctx, dec)
+
+	// Compute the start offset of the payload in msg.
+	payloadOffset := dec.Int() + 8 /*int64 that stores the length of the header*/


Need to check that the decoded number is correct.

The more I look at it, the more I think that we should not try to store the encoded header length inside the encoded header. Instead we can do:

header_length : uint32 header : byte[header_length] ... rest is payload ...

So we read the first 4 bytes to get the header length, use a Decoder to parse just the header, and leave the rest as the payload.

Changed the encoding based on your suggestions.

ghemawat · 2024-08-06T15:53:06Z

internal/net/call/call.go

 	"github.com/ServiceWeaver/weaver/runtime/logging"
 	"github.com/ServiceWeaver/weaver/runtime/retry"
 	"go.opentelemetry.io/otel/codes"
 	"go.opentelemetry.io/otel/trace"
 )

-const (


There is a long comment in msg.go that explains the message format. That should be updated, but I don't see it in this commit.

Good point. Updated.

metadata/metadata.go

ghemawat · 2024-08-06T15:58:40Z

metadata/metadata.go

+	if !ok {
+		return nil, false
+	}
+	out := make(map[string]string, len(meta))


Can we use maps.Clone instead?

Good point. Done

ghemawat · 2024-08-06T15:59:32Z

weavertest/internal/simple/simple.go

@@ -28,6 +28,8 @@ import (
 	"sync"

 	"github.com/ServiceWeaver/weaver"
+	"github.com/ServiceWeaver/weaver/metadata"
+	"golang.org/x/exp/maps"


Do we still need x/exp/maps, or can we use the maps package in the standard library now?

No, replaced with "maps". Thank you

rgrandl

Thanks, Sanjay

rgrandl · 2024-08-07T17:46:00Z

internal/net/call/call.go

-	var hdr [msgHeaderSize]byte
-	copy(hdr[0:], h[:])
+	enc := codegen.NewEncoder()
+	enc.Bytes(h[:])


Good point. Changed.

rgrandl · 2024-08-07T17:46:28Z

internal/net/call/call.go

+	// Encode the length of the header. Note that the header can have variable length,
+	// so we need to know at which offset in the encoded message starts the payload
+	// when we decode the message at the receiver.
+	enc.Int(len(enc.Data()))


We don't need to encode the length of the header based on your suggestions.

rgrandl · 2024-08-07T17:46:58Z

internal/net/call/call.go

+	// Extract trace context and create a new child span to trace the method
+	// call on the server.
+	span := trace.SpanFromContext(ctx) // noop span
+	if sc := readTraceContext(dec); sc != nil && sc.IsValid() {


Good point. We are not failing the RPC right now if sc is invalid, but I think we should do that. Changed.

rgrandl · 2024-08-07T17:47:18Z

internal/net/call/call.go

+	ctx = readContextMetadata(ctx, dec)
+
+	// Compute the start offset of the payload in msg.
+	payloadOffset := dec.Int() + 8 /*int64 that stores the length of the header*/


Changed the encoding based on your suggestions.

metadata/metadata.go

rgrandl · 2024-08-07T17:48:13Z

metadata/metadata.go

+	if !ok {
+		return nil, false
+	}
+	out := make(map[string]string, len(meta))


Good point. Done

rgrandl · 2024-08-07T17:48:28Z

weavertest/internal/simple/simple.go

@@ -28,6 +28,8 @@ import (
 	"sync"

 	"github.com/ServiceWeaver/weaver"
+	"github.com/ServiceWeaver/weaver/metadata"
+	"golang.org/x/exp/maps"


No, replaced with "maps". Thank you

rgrandl · 2024-08-07T17:57:51Z

internal/net/call/call.go

 	"github.com/ServiceWeaver/weaver/runtime/logging"
 	"github.com/ServiceWeaver/weaver/runtime/retry"
 	"go.opentelemetry.io/otel/codes"
 	"go.opentelemetry.io/otel/trace"
 )

-const (


Good point. Updated.

rgrandl · 2024-08-07T18:25:37Z

Sanjay, don't look at the PR yet. Some test is failing.

In the current implementation, weaver doesn't allow the user to propagate context information. We recommend users to define a struct that encapsulates the metadata information and add it as an argument to the method. However, more and more users are asking for an option to propagate metadata information using the context. This request comes especially from users that are using gRPC to communicate between their services, and gRPC provides a way to propagate metadata information using the context. This PR enables the users to propagate metadata information as a map[string]string. ```main.go // To attach metadata with key "foo" and value "bar" to the context, you can do: ctx := context.Background() ctx = metadata.NewContext(ctx, map[string]string{"foo": "bar"}) // To read the metadata value associated with a key "foo" in the context, you can do: meta, found := metadata.FromContext(ctx) if found { value := meta["foo"] } ``` [1] https://pkg.go.dev/google.golang.org/grpc/metadata

rgrandl · 2024-08-07T19:45:25Z

Sanjay, don't look at the PR yet. Some test is failing.

Never mind. Everything works as expected. We have a flaky test but that's not something introduced by this change.

ghemawat · 2024-08-07T20:25:16Z

internal/net/call/call.go

+
+	// Note that we send the header and the payload as follows:
+	// [header_length][encoded_header][payload]
+	hdrSlice := make([]byte, 4)


Unnecessary allocation. You can do:

var hdr [4]byte binary...PutUint32(hdr[:], ...)

ghemawat · 2024-08-07T20:25:50Z

internal/net/call/call.go

+	// Note that we send the header and the payload as follows:
+	// [header_length][encoded_header][payload]
+	hdrSlice := make([]byte, 4)
+	binary.BigEndian.PutUint32(hdrSlice, uint32(len(enc.Data())))


Why BigEndian? I think we are using LittleEndian everywhere else.

Good point.

ghemawat · 2024-08-07T20:27:26Z

internal/net/call/call.go

-	// Extract request header from front of payload.
-	if len(msg) < msgHeaderSize {
+	msgLen := uint32(len(msg))
+	hdrLenEndOffset := uint32(4)


This can be a "const hdrLenEndOffset".

Also, maybe hdrLenLen would be a better name?

ghemawat · 2024-08-07T20:28:33Z

internal/net/call/call.go

+	// call on the server.
+	span := trace.SpanFromContext(ctx) // noop span
+	if sc := readTraceContext(dec); sc != nil {
+		if sc.IsValid() {


How about checking for error and returning early on error instead of placing the error-handling in the else branch?

ghemawat · 2024-08-07T20:29:40Z

internal/net/call/msg.go

-//    traceContext [25]byte   -- zero, or trace context
-//    remainder               -- call argument serialization
+//    headerLen         [4]byte       -- length of the encoded header
+//    header            [length]byte  -- encoded header information


Reword the type of header of [headerLen]byte to match the name used above

ghemawat · 2024-08-07T20:31:19Z

internal/net/call/msg.go

-//    remainder               -- call argument serialization
+//    headerLen         [4]byte       -- length of the encoded header
+//    header            [length]byte  -- encoded header information
+//      headerKey       [16]byte      -- fingerprint of method name


The following part looks wrong? I think you need follow header by remainder and then extend the comment after that to say something like:

The header is encoded using Service Weaver's encoding format for a type that looks like:

struct { ...fill this in... }

Good idea. Done

ghemawat · 2024-08-07T20:32:02Z

internal/net/call/msg.go

+//      deadline        [8]byte       -- zero, or deadline in microseconds
+//      traceContext    [25]byte      -- zero, or trace context
+//      metadataContext [length]byte  -- encoded map[string]string
+//    remainder                       -- call argument serialization


payload would be a better name than remainder

ghemawat · 2024-08-07T20:34:10Z

weavertest/internal/simple/simple.go

+	d.mu.Lock()
+	defer d.mu.Unlock()
+	meta, found := metadata.FromContext(ctx)
+	if found {


Can do:

if meta, found := ...; found {

ghemawat · 2024-08-07T20:35:43Z

internal/net/call/call.go

@@ -385,25 +381,36 @@ func (rc *reconnectingConnection) Call(ctx context.Context, h MethodKey, arg []b
 }

 func (rc *reconnectingConnection) callOnce(ctx context.Context, h MethodKey, arg []byte, opts CallOptions) ([]byte, error) {
-	var hdr [msgHeaderSize]byte
-	copy(hdr[0:], h[:])
+	enc := codegen.NewEncoder()


How about moving the header encoding/decoding code into routines of their own. The code here can then be:

hdr := encodeHeader(ctx, h)

rgrandl · 2024-08-07T22:47:10Z

Thanks Sanjay

ghemawat · 2024-08-07T22:48:51Z

internal/net/call/call.go

+	var hdrLen [hdrLenLen]byte
+	binary.LittleEndian.PutUint32(hdrLen[:], uint32(len(hdr)))
+	hdrSlice := hdrLen[:]
+	hdrSlice = append(hdrSlice, hdr...)


Maybe just combine prev 2 lines?
hdrSlice := append(hdrLen[:], hdr...)

ghemawat · 2024-08-07T22:51:04Z

internal/net/call/msg.go

+// The header is encoded using Service Weaver's encoding format for a type that
+// looks like:
+//
+// struct header {


Line up the types just like "go fmt" would.

rgrandl

Thanks Sanjay. Looks great after your suggestions

rgrandl · 2024-08-07T22:57:04Z

internal/net/call/call.go

+	var hdrLen [hdrLenLen]byte
+	binary.LittleEndian.PutUint32(hdrLen[:], uint32(len(hdr)))
+	hdrSlice := hdrLen[:]
+	hdrSlice = append(hdrSlice, hdr...)


rgrandl · 2024-08-07T22:58:16Z

internal/net/call/msg.go

+// The header is encoded using Service Weaver's encoding format for a type that
+// looks like:
+//
+// struct header {


rgrandl requested a review from ghemawat July 17, 2024 20:52

rgrandl force-pushed the context_propagation branch 2 times, most recently from a91af38 to 8225fbf Compare July 17, 2024 23:40

rgrandl commented Jul 17, 2024

View reviewed changes

rgrandl self-assigned this Jul 17, 2024

ghemawat reviewed Jul 18, 2024

View reviewed changes

rgrandl force-pushed the context_propagation branch 2 times, most recently from a6d09c3 to b715910 Compare August 5, 2024 23:43

rgrandl commented Aug 5, 2024

View reviewed changes

ghemawat reviewed Aug 6, 2024

View reviewed changes

rgrandl force-pushed the context_propagation branch from b715910 to d478770 Compare August 7, 2024 17:58

rgrandl commented Aug 7, 2024

View reviewed changes

rgrandl force-pushed the context_propagation branch from d478770 to 492c104 Compare August 7, 2024 19:36

ghemawat reviewed Aug 7, 2024

View reviewed changes

Header enc/dec into their own routines

8e630e4

ghemawat approved these changes Aug 7, 2024

View reviewed changes

Last fixes

344e885

rgrandl commented Aug 7, 2024

View reviewed changes

rgrandl merged commit d613ffe into main Aug 7, 2024
10 checks passed

rgrandl deleted the context_propagation branch August 7, 2024 23:54

Add context propagation support #783

Add context propagation support #783

Conversation

rgrandl commented Jul 17, 2024

rgrandl commented Jul 17, 2024

rgrandl left a comment

Choose a reason for hiding this comment

ghemawat left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

rgrandl left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

rgrandl left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

rgrandl commented Aug 7, 2024

rgrandl commented Aug 7, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

rgrandl commented Aug 7, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment