diff --git a/src/x/serialize/encoder.go b/src/x/serialize/encoder.go index c54859ea2e..63b94a6191 100644 --- a/src/x/serialize/encoder.go +++ b/src/x/serialize/encoder.go @@ -68,6 +68,11 @@ var defaultNewCheckedBytesFn = checked.NewBytes type encoder struct { buf *bytes.Buffer checkedBytes checked.Bytes + numTags uint16 + + customTag []byte + customStart int + customEnd int opts TagEncoderOptions pool TagEncoderPool @@ -77,6 +82,15 @@ func newTagEncoder( newFn newCheckedBytesFn, opts TagEncoderOptions, pool TagEncoderPool, +) TagEncoder { + return newTagEncoderWithCustomTag(newFn, opts, pool, nil) +} + +func newTagEncoderWithCustomTag( + newFn newCheckedBytesFn, + opts TagEncoderOptions, + pool TagEncoderPool, + customTag []byte, ) TagEncoder { b := make([]byte, 0, opts.InitialCapacity()) cb := newFn(nil, nil) @@ -85,6 +99,9 @@ func newTagEncoder( checkedBytes: cb, opts: opts, pool: pool, + customStart: -1, + customEnd: -1, + customTag: customTag, } } @@ -97,6 +114,7 @@ func (e *encoder) Encode(srcTags ident.TagIterator) error { defer tags.Close() numTags := tags.Remaining() + e.numTags = uint16(numTags) max := int(e.opts.TagSerializationLimits().MaxNumberTags()) if numTags > max { return fmt.Errorf("too many tags to encode (%d), limit is: %d", numTags, max) @@ -138,6 +156,23 @@ func (e *encoder) Data() (checked.Bytes, bool) { return e.checkedBytes, true } +func (e *encoder) ResetCustomTag() { + if e.checkedBytes.NumRef() == 0 { + return + } + if e.customStart != -1 { + b := e.checkedBytes.Bytes() + b = append(b[:e.customStart], b[e.customEnd:]...) + + // Also need to decrement the number of tags being encoded. + e.numTags-- + numTags := encodeUInt16(e.numTags) + b[len(headerMagicBytes)] = numTags[0] + b[len(headerMagicBytes)+1] = numTags[1] + e.checkedBytes.Reset(b) + } +} + func (e *encoder) Reset() { if e.checkedBytes.NumRef() == 0 { return @@ -145,6 +180,8 @@ func (e *encoder) Reset() { e.buf.Reset() e.checkedBytes.Reset(nil) e.checkedBytes.DecRef() + e.customStart = -1 + e.customEnd = -1 } func (e *encoder) Finalize() { @@ -157,6 +194,18 @@ func (e *encoder) Finalize() { } func (e *encoder) encodeTag(t ident.Tag) error { + // If this is the custom tag mark the start and end of the tag encoding, + // so that it maybe removed later. + var ( + start int + customTag bool + ) + if len(e.customTag) > 0 && + bytes.Equal(t.Name.Bytes(), e.customTag) { + start = e.buf.Len() + customTag = true + } + if len(t.Name.Bytes()) == 0 { return errEmptyTagNameLiteral } @@ -165,7 +214,16 @@ func (e *encoder) encodeTag(t ident.Tag) error { return err } - return e.encodeID(t.Value) + if err := e.encodeID(t.Value); err != nil { + return err + } + + // If this was the custom tag, then store its encoding information. + if customTag { + e.customStart = start + e.customEnd = e.buf.Len() + } + return nil } func (e *encoder) encodeID(i ident.ID) error { diff --git a/src/x/serialize/encoder_test.go b/src/x/serialize/encoder_test.go index fe4d81651d..35b250f8a5 100644 --- a/src/x/serialize/encoder_test.go +++ b/src/x/serialize/encoder_test.go @@ -119,6 +119,66 @@ func TestSimpleEncode(t *testing.T) { require.Equal(t, "bar", string(b[20:23])) } +func TestEncodeCustomTagReset(t *testing.T) { + e := newTagEncoderWithCustomTag(defaultNewCheckedBytesFn, newTestEncoderOpts(), nil, []byte("__m3_type__")) + + tags := ident.NewTagsIterator(ident.NewTags( + ident.StringTag("abc", "defg"), + ident.StringTag("__m3_type__", "gauge"), + ident.StringTag("x", "bar"), + )) + require.NoError(t, e.Encode(tags)) + + bc, ok := e.Data() + require.True(t, ok) + require.NotNil(t, bc) + b := bc.Bytes() + numExpectedBytes := 2 /* header */ + 2 /* num tags */ + + 2 /* abc length */ + len("abc") + + 2 /* defg length */ + len("defg") + + 2 /* x length */ + len("__m3_type__") + + 2 /* bar length */ + len("gauge") + + 2 /* x length */ + len("x") + + 2 /* bar length */ + len("bar") + require.Len(t, b, numExpectedBytes) + require.Equal(t, headerMagicBytes, b[:2]) + require.Equal(t, encodeUInt16(3), b[2:4]) + require.Equal(t, uint16(3), decodeUInt16(b[4:6])) /* len abc */ + require.Equal(t, "abc", string(b[6:9])) + require.Equal(t, uint16(4), decodeUInt16(b[9:11])) /* len defg */ + require.Equal(t, "defg", string(b[11:15])) + require.Equal(t, uint16(11), decodeUInt16(b[15:17])) /* len __m3_type__ */ + require.Equal(t, "__m3_type__", string(b[17:28])) + require.Equal(t, uint16(5), decodeUInt16(b[28:30])) /* len gauge */ + require.Equal(t, "gauge", string(b[30:35])) + require.Equal(t, uint16(1), decodeUInt16(b[35:37])) /* len x */ + require.Equal(t, "x", string(b[37:38])) + require.Equal(t, uint16(3), decodeUInt16(b[38:40])) /* len bar */ + require.Equal(t, "bar", string(b[40:43])) + + e.ResetCustomTag() + bc, ok = e.Data() + require.True(t, ok) + require.NotNil(t, bc) + b = bc.Bytes() + numExpectedBytes = 2 /* header */ + 2 /* num tags */ + + 2 /* abc length */ + len("abc") + + 2 /* defg length */ + len("defg") + + 2 /* x length */ + len("x") + + 2 /* bar length */ + len("bar") + require.Len(t, b, numExpectedBytes) + require.Equal(t, headerMagicBytes, b[:2]) + require.Equal(t, encodeUInt16(2), b[2:4]) + require.Equal(t, uint16(3), decodeUInt16(b[4:6])) /* len abc */ + require.Equal(t, "abc", string(b[6:9])) + require.Equal(t, uint16(4), decodeUInt16(b[9:11])) /* len defg */ + require.Equal(t, "defg", string(b[11:15])) + require.Equal(t, uint16(1), decodeUInt16(b[15:17])) /* len x */ + require.Equal(t, "x", string(b[17:18])) + require.Equal(t, uint16(3), decodeUInt16(b[18:20])) /* len bar */ + require.Equal(t, "bar", string(b[20:23])) +} + func TestTagEncoderErrorEncoding(t *testing.T) { opts := NewTagEncoderOptions() e := newTagEncoder(defaultNewCheckedBytesFn, opts, nil) diff --git a/src/x/serialize/types.go b/src/x/serialize/types.go index a694ae7dc7..ef04e01a73 100644 --- a/src/x/serialize/types.go +++ b/src/x/serialize/types.go @@ -45,6 +45,9 @@ type TagEncoder interface { // TagEncoder. Data() (checked.Bytes, bool) + // ResetCustomTag removes the custom tag from the completed byte encoding. + ResetCustomTag() + // Reset resets the internal state to allow reuse of the encoder. Reset()