Repository: twmb/franz-go
Pull request #255 (2022-11-20 03:22:23)

pkg/sr: pin to go1.19, convert to `any`, add support for indices
We pin to 1.19 so that we can pull in the new binary.AppendVarint, and in doing so, we convert to `any` for simplicity.

This adds support for a type to be indexes, which means the registry now supports protobuf definitions. We also add a few helper APIs.

```
    func GenerateFn(func() any) SerdeOpt
    func Index(...int) SerdeOpt
    func (*Serde) SetDefaults(...SerdeOpt)
    func (*Serde) DecodeNew([]byte) any
```

For #207


Patch:
diff --git a/pkg/sr/client.go b/pkg/sr/client.go
index 2657a8bb..5583f111 100644
--- a/pkg/sr/client.go
+++ b/pkg/sr/client.go
@@ -77,23 +77,23 @@ func NewClient(opts ...Opt) (*Client, error) {
 	return cl, nil
 }
 
-func (cl *Client) get(ctx context.Context, path string, into interface{}) error {
+func (cl *Client) get(ctx context.Context, path string, into any) error {
 	return cl.do(ctx, http.MethodGet, path, nil, into)
 }
 
-func (cl *Client) post(ctx context.Context, path string, v, into interface{}) error {
+func (cl *Client) post(ctx context.Context, path string, v, into any) error {
 	return cl.do(ctx, http.MethodPost, path, v, into)
 }
 
-func (cl *Client) put(ctx context.Context, path string, v, into interface{}) error {
+func (cl *Client) put(ctx context.Context, path string, v, into any) error {
 	return cl.do(ctx, http.MethodPut, path, v, into)
 }
 
-func (cl *Client) delete(ctx context.Context, path string, into interface{}) error {
+func (cl *Client) delete(ctx context.Context, path string, into any) error {
 	return cl.do(ctx, http.MethodDelete, path, nil, into)
 }
 
-func (cl *Client) do(ctx context.Context, method, path string, v, into interface{}) error {
+func (cl *Client) do(ctx context.Context, method, path string, v, into any) error {
 	urls := cl.urls
 
 start:
diff --git a/pkg/sr/go.mod b/pkg/sr/go.mod
index 7496fd7f..cdc6b1a8 100644
--- a/pkg/sr/go.mod
+++ b/pkg/sr/go.mod
@@ -1,3 +1,3 @@
 module github.com/twmb/franz-go/pkg/sr
 
-go 1.15
+go 1.19
diff --git a/pkg/sr/serde.go b/pkg/sr/serde.go
index 4c7e5f3c..6445fff8 100644
--- a/pkg/sr/serde.go
+++ b/pkg/sr/serde.go
@@ -3,6 +3,7 @@ package sr
 import (
 	"encoding/binary"
 	"errors"
+	"io"
 	"reflect"
 	"sync"
 	"sync/atomic"
@@ -33,26 +34,54 @@ type (
 func (o serdeOpt) apply(t *tserde) { o.fn(t) }
 
 // EncodeFn allows Serde to encode a value.
-func EncodeFn(fn func(interface{}) ([]byte, error)) SerdeOpt {
+func EncodeFn(fn func(any) ([]byte, error)) SerdeOpt {
 	return serdeOpt{func(t *tserde) { t.encode = fn }}
 }
 
 // AppendEncodeFn allows Serde to encode a value to an existing slice. This
 // can be more efficient than EncodeFn; this function is used if it exists.
-func AppendEncodeFn(fn func([]byte, interface{}) ([]byte, error)) SerdeOpt {
+func AppendEncodeFn(fn func([]byte, any) ([]byte, error)) SerdeOpt {
 	return serdeOpt{func(t *tserde) { t.appendEncode = fn }}
 }
 
 // DecodeFn allows Serde to decode into a value.
-func DecodeFn(fn func([]byte, interface{}) error) SerdeOpt {
+func DecodeFn(fn func([]byte, any) error) SerdeOpt {
 	return serdeOpt{func(t *tserde) { t.decode = fn }}
 }
 
+// GenerateFn returns a new(Value) that can be decoded into. This function can
+// be used to control the instantiation of a new type for DecodeNew.
+func GenerateFn(fn func() any) SerdeOpt {
+	return serdeOpt{func(t *tserde) { t.gen = fn }}
+}
+
+// Index attaches a message index to a value. A single schema ID can be
+// registered multiple times with different indices.
+//
+// This option supports schemas that encode many different values from the same
+// schema (namely, protobuf). The index into the the schema to encode a
+// particular message is specified with `index`.
+//
+// NOTE: this option must be used for protobuf schemas.
+//
+// For more information, see where `message-indexes` are described in:
+//
+//	https://docs.confluent.io/platform/current/schema-registry/serdes-develop/index.html#wire-format
+func Index(index ...int) SerdeOpt {
+	return serdeOpt{func(t *tserde) { t.index = index }}
+}
+
 type tserde struct {
 	id           uint32
-	encode       func(interface{}) ([]byte, error)
-	appendEncode func([]byte, interface{}) ([]byte, error)
-	decode       func([]byte, interface{}) error
+	exists       bool
+	encode       func(any) ([]byte, error)
+	appendEncode func([]byte, any) ([]byte, error)
+	decode       func([]byte, any) error
+	gen          func() any
+	typeof       reflect.Type
+
+	index    []int          // for encoding, an optional index we use
+	subindex map[int]tserde // for decoding, we look up sub-indices in the payload
 }
 
 // Serde encodes and decodes values according to the schema registry wire
@@ -70,6 +99,8 @@ type Serde struct {
 	ids   atomic.Value // map[int]tserde
 	types atomic.Value // map[reflect.Type]tserde
 	mu    sync.Mutex
+
+	defaults []SerdeOpt
 }
 
 var (
@@ -93,47 +124,111 @@ func (s *Serde) loadTypes() map[reflect.Type]tserde {
 	return types.(map[reflect.Type]tserde)
 }
 
+// SetDefaults sets default options to apply to every registered type. These
+// options are always applied first, so you can override them as necessary when
+// registering.
+//
+// This can be useful if you always want to use the same encoding or decoding
+// functions.
+func (s *Serde) SetDefaults(opts ...SerdeOpt) {
+	s.defaults = opts
+}
+
 // Register registers a schema ID and the value it corresponds to, as well as
 // the encoding or decoding functions. You need to register functions depending
 // on whether you are only encoding, only decoding, or both.
-func (s *Serde) Register(id int, v interface{}, opts ...SerdeOpt) {
-	t := tserde{id: uint32(id)}
+func (s *Serde) Register(id int, v any, opts ...SerdeOpt) {
+	var t tserde
+	for _, opt := range s.defaults {
+		opt.apply(&t)
+	}
 	for _, opt := range opts {
 		opt.apply(&t)
 	}
 
+	typeof := reflect.TypeOf(v)
+
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
-	{
-		dup := make(map[int]tserde)
-		for k, v := range s.loadIDs() {
-			dup[k] = v
-		}
-		dup[id] = t
-		s.ids.Store(dup)
+	// Type mapping is easy: we just map the type to the final tserde.
+	// We defer the store because we modify the tserde below, and we
+	// may delete a type key.
+	dupTypes := make(map[reflect.Type]tserde)
+	for k, v := range s.loadTypes() {
+		dupTypes[k] = v
 	}
+	defer func() {
+		dupTypes[typeof] = t
+		s.types.Store(dupTypes)
+	}()
 
-	{
-		dup := make(map[reflect.Type]tserde)
-		for k, v := range s.loadTypes() {
-			dup[k] = v
-		}
-		dup[reflect.TypeOf(v)] = t
-		s.types.Store(dup)
+	// For IDs, we deeply clone any path that is changing.
+	m := tserdeMapClone(s.loadIDs(), id, t.index)
+	s.ids.Store(m)
+
+	// Now we have a full path down index initialized (or, the top
+	// level map if there is no index). We iterate down the index
+	// tree to find the end node we are initializing.
+	k := id
+	at := m[k]
+	for _, idx := range t.index {
+		m = at.subindex
+		k = idx
+		at = m[k]
 	}
+
+	// If this value was already registered, we are overriding the
+	// previous registration and need to delete the previous type.
+	if at.exists {
+		delete(dupTypes, at.typeof)
+	}
+
+	// Now, we initialize the end node.
+	t = tserde{
+		id:           uint32(id),
+		exists:       true,
+		encode:       t.encode,
+		appendEncode: t.appendEncode,
+		decode:       t.decode,
+		gen:          t.gen,
+		typeof:       typeof,
+		index:        t.index,
+		subindex:     at.subindex,
+	}
+	m[k] = t
+}
+
+func tserdeMapClone(m map[int]tserde, at int, index []int) map[int]tserde {
+	// We deeply clone the input map (or initialize it).
+	dup := make(map[int]tserde, len(m))
+	for k, v := range m {
+		dup[k] = v
+	}
+
+	// If there is an index, we want to ensure the full path to the index
+	// is initialized. Our input "at" is what will get us *to* this index,
+	// and then we need to recurse down index. The caller then maps down
+	// our fully initialized subindex tree and ensures that the terminal
+	// tserde (the end of index) is initialized.
+	if len(index) > 0 {
+		sub := dup[at]
+		sub.subindex = tserdeMapClone(sub.subindex, index[0], index[1:])
+		dup[at] = sub
+	}
+	return dup
 }
 
 // Encode encodes a value according to the schema registry wire format and
 // returns it. If EncodeFn was not used, this returns ErrNotRegistered.
-func (s *Serde) Encode(v interface{}) ([]byte, error) {
+func (s *Serde) Encode(v any) ([]byte, error) {
 	return s.AppendEncode(nil, v)
 }
 
 // AppendEncode appends an encoded value to b according to the schema registry
 // wire format and returns it. If EncodeFn was not used, this returns
 // ErrNotRegistered.
-func (s *Serde) AppendEncode(b []byte, v interface{}) ([]byte, error) {
+func (s *Serde) AppendEncode(b []byte, v any) ([]byte, error) {
 	t, ok := s.loadTypes()[reflect.TypeOf(v)]
 	if !ok || (t.encode == nil && t.appendEncode == nil) {
 		return b, ErrNotRegistered
@@ -147,6 +242,17 @@ func (s *Serde) AppendEncode(b []byte, v interface{}) ([]byte, error) {
 		byte(t.id>>0),
 	)
 
+	if len(t.index) > 0 {
+		if len(t.index) == 1 && t.index[0] == 0 {
+			b = append(b, 0) // first-index shortcut (one type in the protobuf)
+		} else {
+			b = binary.AppendVarint(b, int64(len(t.index)))
+			for _, idx := range t.index {
+				b = binary.AppendVarint(b, int64(idx))
+			}
+		}
+	}
+
 	if t.appendEncode != nil {
 		return t.appendEncode(b, v)
 	}
@@ -159,7 +265,7 @@ func (s *Serde) AppendEncode(b []byte, v interface{}) ([]byte, error) {
 
 // MustEncode returns the value of Encode, panicking on error. This is a
 // shortcut for if your encode function cannot error.
-func (s *Serde) MustEncode(v interface{}) []byte {
+func (s *Serde) MustEncode(v any) []byte {
 	b, err := s.Encode(v)
 	if err != nil {
 		panic(err)
@@ -169,7 +275,7 @@ func (s *Serde) MustEncode(v interface{}) []byte {
 
 // MustAppendEncode returns the value of AppendEncode, panicking on error.
 // This is a shortcut for if your encode function cannot error.
-func (s *Serde) MustAppendEncode(b []byte, v interface{}) []byte {
+func (s *Serde) MustAppendEncode(b []byte, v any) []byte {
 	b, err := s.AppendEncode(b, v)
 	if err != nil {
 		panic(err)
@@ -183,15 +289,71 @@ func (s *Serde) MustAppendEncode(b []byte, v interface{}) []byte {
 // Serde does not handle references in schemas; it is up to you to register the
 // full decode function for any top-level ID, regardless of how many other
 // schemas are referenced in top-level ID.
-func (s *Serde) Decode(b []byte, v interface{}) error {
+func (s *Serde) Decode(b []byte, v any) error {
+	b, t, err := s.decodeFind(b)
+	if err != nil {
+		return err
+	}
+	return t.decode(b, v)
+}
+
+// DecodeNew is the same as Decode, but decodes into a new value rather than
+// the input value. If DecodeFn was not used, this returns ErrNotRegistered.
+// GenerateFn can be used to control the instantiation of a new value,
+// otherwise this uses reflect.New(reflect.TypeOf(v)).Interface().
+func (s *Serde) DecodeNew(b []byte) (any, error) {
+	b, t, err := s.decodeFind(b)
+	if err != nil {
+		return nil, err
+	}
+	var v any
+	if t.gen != nil {
+		v = t.gen()
+	} else {
+		v = reflect.New(t.typeof).Interface()
+	}
+	return v, t.decode(b, v)
+}
+
+func (s *Serde) decodeFind(b []byte) ([]byte, tserde, error) {
 	if len(b) < 5 || b[0] != 0 {
-		return ErrBadHeader
+		return nil, tserde{}, ErrBadHeader
 	}
 	id := binary.BigEndian.Uint32(b[1:5])
+	b = b[5:]
+
+	t := s.loadIDs()[int(id)]
+	if len(t.subindex) > 0 {
+		r := bReader{b}
+		br := io.ByteReader(&r)
+		l, err := binary.ReadVarint(br)
+		if l == 0 { // length 0 is a shortcut for length 1, index 0
+			t = t.subindex[0]
+		}
+		for err == nil && t.subindex != nil && l > 0 {
+			var idx int64
+			idx, err = binary.ReadVarint(br)
+			t = t.subindex[int(idx)]
+			l--
+		}
+		if err != nil {
+			return nil, t, err
+		}
+		b = r.b
+	}
+	if !t.exists {
+		return nil, t, ErrNotRegistered
+	}
+	return b, t, nil
+}
+
+type bReader struct{ b []byte }
 
-	t, ok := s.loadIDs()[int(id)]
-	if !ok || t.decode == nil {
-		return ErrNotRegistered
+func (b *bReader) ReadByte() (byte, error) {
+	if len(b.b) > 0 {
+		r := b.b[0]
+		b.b = b.b[1:]
+		return r, nil
 	}
-	return t.decode(b[5:], v)
+	return 0, io.EOF
 }


Test patch:
diff --git a/pkg/sr/serde_test.go b/pkg/sr/serde_test.go
new file mode 100644
index 00000000..8667a0d2
--- /dev/null
+++ b/pkg/sr/serde_test.go
@@ -0,0 +1,141 @@
+package sr
+
+import (
+	"bytes"
+	"encoding/json"
+	"io"
+	"reflect"
+	"testing"
+)
+
+func TestSerde(t *testing.T) {
+	type (
+		overridden struct {
+			Ignored string `json:"ignored,omitempty"`
+		}
+		overrides struct {
+			One string `json:"one,omitempty"`
+		}
+		idx1 struct {
+			Two   int64 `json:"two,omitempty"`
+			Three int64 `json:"three,omitempty"`
+		}
+		idx2 struct {
+			Biz string `json:"biz,omitempty"`
+			Baz string `json:"baz,omitempty"`
+		}
+		idx3 struct {
+			Boz int8 `json:"boz,omitempty"`
+		}
+		idx4 struct {
+			Bingo string `json:"bingo,omitempty"`
+		}
+		oneidx struct {
+			Foo string `json:"foo,omitempty"`
+			Bar string `json:"bar,omitempty"`
+		}
+	)
+
+	var serde Serde
+	serde.SetDefaults(
+		EncodeFn(json.Marshal),
+		DecodeFn(json.Unmarshal),
+	)
+	serde.Register(127, overridden{}, GenerateFn(func() any { return new(overridden) }))
+	serde.Register(127, overrides{})
+	serde.Register(3, idx1{}, Index(0))
+	serde.Register(3, idx2{}, Index(1), AppendEncodeFn(func(b []byte, v any) ([]byte, error) {
+		bb, err := json.Marshal(v)
+		if err != nil {
+			return b, err
+		}
+		return append(b, bb...), nil
+	}))
+	serde.Register(3, idx4{}, Index(0, 0, 1))
+	serde.Register(3, idx3{}, Index(0, 0))
+	serde.Register(5, oneidx{}, Index(0), GenerateFn(func() any { return &oneidx{Foo: "defoo", Bar: "debar"} }))
+
+	for i, test := range []struct {
+		enc    any
+		expEnc []byte
+		expDec any
+		expErr bool
+	}{
+		{
+			enc:    overridden{},
+			expErr: true,
+		},
+		{
+			enc:    overrides{"foo"},
+			expEnc: append([]byte{0, 0, 0, 0, 127}, `{"one":"foo"}`...),
+		},
+		{
+			enc:    idx1{Two: 2, Three: 3},
+			expEnc: append([]byte{0, 0, 0, 0, 3, 0}, `{"two":2,"three":3}`...),
+		},
+		{
+			enc:    idx2{Biz: "bizzy", Baz: "bazzy"},
+			expEnc: append([]byte{0, 0, 0, 0, 3, 2, 2}, `{"biz":"bizzy","baz":"bazzy"}`...),
+		},
+		{
+			enc:    idx3{Boz: 8},
+			expEnc: append([]byte{0, 0, 0, 0, 3, 4, 0, 0}, `{"boz":8}`...),
+		},
+		{
+			enc:    idx4{Bingo: "bango"},
+			expEnc: append([]byte{0, 0, 0, 0, 3, 6, 0, 0, 2}, `{"bingo":"bango"}`...),
+		},
+		{
+			enc:    oneidx{Bar: "bar"},
+			expEnc: append([]byte{0, 0, 0, 0, 5, 0}, `{"bar":"bar"}`...),
+			expDec: oneidx{Foo: "defoo", Bar: "bar"},
+		},
+	} {
+		b, err := serde.Encode(test.enc)
+		gotErr := err != nil
+		if gotErr != test.expErr {
+			t.Errorf("#%d Encode: got err? %v, exp err? %v", i, gotErr, test.expErr)
+			continue
+		}
+		if test.expErr {
+			continue
+		}
+
+		if !bytes.Equal(b, test.expEnc) {
+			t.Errorf("#%d: Encode(%v) != exp(%v)", i, b, test.expEnc)
+			continue
+		}
+
+		if b2 := serde.MustEncode(test.enc); !bytes.Equal(b2, b) {
+			t.Errorf("#%d got MustEncode(%v) != Encode(%v)", i, b2, b)
+		}
+		if b2 := serde.MustAppendEncode([]byte("foo"), test.enc); !bytes.Equal(b2, append([]byte("foo"), b...)) {
+			t.Errorf("#%d got MustAppendEncode(%v) != Encode(foo%v)", i, b2, b)
+		}
+
+		v, err := serde.DecodeNew(b)
+		if err != nil {
+			t.Errorf("#%d DecodeNew: got unexpected err %v", i, err)
+			continue
+		}
+		v = reflect.Indirect(reflect.ValueOf(v)).Interface() // DecodeNew returns a pointer, we compare values below
+
+		exp := test.expDec
+		if exp == nil {
+			exp = test.enc
+		}
+		if !reflect.DeepEqual(v, exp) {
+			t.Errorf("#%d round trip: got %v != exp %v", i, v, exp)
+		}
+	}
+
+	if _, err := serde.DecodeNew([]byte{1, 0, 0, 0, 0, 0}); err != ErrBadHeader {
+		t.Errorf("got %v != exp ErrBadHeader", err)
+	}
+	if _, err := serde.DecodeNew([]byte{0, 0, 0, 0, 3}); err != io.EOF {
+		t.Errorf("got %v != exp io.EOF", err)
+	}
+	if _, err := serde.DecodeNew([]byte{0, 0, 0, 0, 99}); err != ErrNotRegistered {
+		t.Errorf("got %v != exp ErrNotRegistered", err)
+	}
+}