Implement the marshaler interfaces for Rule/RuleOP

Signed-off-by: JmPotato <[email protected]>
tikv · Nov 28, 2023 · eeaf477 · eeaf477
1 parent a6e855e
commit eeaf477
Show file tree

Hide file tree

Showing 3 changed files with 239 additions and 0 deletions.
diff --git a/client/http/codec.go b/client/http/codec.go
@@ -0,0 +1,99 @@
+// Copyright 2023 TiKV Project Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package http
+
+import (
+	"github.com/pingcap/errors"
+)
+
+const (
+	encGroupSize = 8
+	encMarker    = byte(0xFF)
+	encPad       = byte(0x0)
+)
+
+var pads = make([]byte, encGroupSize)
+
+// encodeBytes guarantees the encoded value is in ascending order for comparison,
+// encoding with the following rule:
+//
+//	[group1][marker1]...[groupN][markerN]
+//	group is 8 bytes slice which is padding with 0.
+//	marker is `0xFF - padding 0 count`
+//
+// For example:
+//
+//	[] -> [0, 0, 0, 0, 0, 0, 0, 0, 247]
+//	[1, 2, 3] -> [1, 2, 3, 0, 0, 0, 0, 0, 250]
+//	[1, 2, 3, 0] -> [1, 2, 3, 0, 0, 0, 0, 0, 251]
+//	[1, 2, 3, 4, 5, 6, 7, 8] -> [1, 2, 3, 4, 5, 6, 7, 8, 255, 0, 0, 0, 0, 0, 0, 0, 0, 247]
+//
+// Refer: https://github.com/facebook/mysql-5.6/wiki/MyRocks-record-format#memcomparable-format
+func encodeBytes(data []byte) []byte {
+	// Allocate more space to avoid unnecessary slice growing.
+	// Assume that the byte slice size is about `(len(data) / encGroupSize + 1) * (encGroupSize + 1)` bytes,
+	// that is `(len(data) / 8 + 1) * 9` in our implement.
+	dLen := len(data)
+	result := make([]byte, 0, (dLen/encGroupSize+1)*(encGroupSize+1))
+	for idx := 0; idx <= dLen; idx += encGroupSize {
+		remain := dLen - idx
+		padCount := 0
+		if remain >= encGroupSize {
+			result = append(result, data[idx:idx+encGroupSize]...)
+		} else {
+			padCount = encGroupSize - remain
+			result = append(result, data[idx:]...)
+			result = append(result, pads[:padCount]...)
+		}
+
+		marker := encMarker - byte(padCount)
+		result = append(result, marker)
+	}
+	return result
+}
+
+func decodeBytes(b []byte) ([]byte, []byte, error) {
+	buf := make([]byte, 0, len(b))
+	for {
+		if len(b) < encGroupSize+1 {
+			return nil, nil, errors.New("insufficient bytes to decode value")
+		}
+
+		groupBytes := b[:encGroupSize+1]
+
+		group := groupBytes[:encGroupSize]
+		marker := groupBytes[encGroupSize]
+
+		padCount := encMarker - marker
+		if padCount > encGroupSize {
+			return nil, nil, errors.Errorf("invalid marker byte, group bytes %q", groupBytes)
+		}
+
+		realGroupSize := encGroupSize - padCount
+		buf = append(buf, group[:realGroupSize]...)
+		b = b[encGroupSize+1:]
+
+		if padCount != 0 {
+			// Check validity of padding bytes.
+			for _, v := range group[realGroupSize:] {
+				if v != encPad {
+					return nil, nil, errors.Errorf("invalid padding byte, group bytes %q", groupBytes)
+				}
+			}
+			break
+		}
+	}
+	return b, buf, nil
+}
diff --git a/client/http/codec_test.go b/client/http/codec_test.go
@@ -0,0 +1,64 @@
+// Copyright 2023 TiKV Project Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package http
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestBytesCodec(t *testing.T) {
+	inputs := []struct {
+		enc []byte
+		dec []byte
+	}{
+		{[]byte{}, []byte{0, 0, 0, 0, 0, 0, 0, 0, 247}},
+		{[]byte{0}, []byte{0, 0, 0, 0, 0, 0, 0, 0, 248}},
+		{[]byte{1, 2, 3}, []byte{1, 2, 3, 0, 0, 0, 0, 0, 250}},
+		{[]byte{1, 2, 3, 0}, []byte{1, 2, 3, 0, 0, 0, 0, 0, 251}},
+		{[]byte{1, 2, 3, 4, 5, 6, 7}, []byte{1, 2, 3, 4, 5, 6, 7, 0, 254}},
+		{[]byte{0, 0, 0, 0, 0, 0, 0, 0}, []byte{0, 0, 0, 0, 0, 0, 0, 0, 255, 0, 0, 0, 0, 0, 0, 0, 0, 247}},
+		{[]byte{1, 2, 3, 4, 5, 6, 7, 8}, []byte{1, 2, 3, 4, 5, 6, 7, 8, 255, 0, 0, 0, 0, 0, 0, 0, 0, 247}},
+		{[]byte{1, 2, 3, 4, 5, 6, 7, 8, 9}, []byte{1, 2, 3, 4, 5, 6, 7, 8, 255, 9, 0, 0, 0, 0, 0, 0, 0, 248}},
+	}
+
+	for _, input := range inputs {
+		b := encodeBytes(input.enc)
+		require.Equal(t, input.dec, b)
+
+		_, d, err := decodeBytes(b)
+		require.NoError(t, err)
+		require.Equal(t, input.enc, d)
+	}
+
+	// Test error decode.
+	errInputs := [][]byte{
+		{1, 2, 3, 4},
+		{0, 0, 0, 0, 0, 0, 0, 247},
+		{0, 0, 0, 0, 0, 0, 0, 0, 246},
+		{0, 0, 0, 0, 0, 0, 0, 1, 247},
+		{1, 2, 3, 4, 5, 6, 7, 8, 0},
+		{1, 2, 3, 4, 5, 6, 7, 8, 255, 1},
+		{1, 2, 3, 4, 5, 6, 7, 8, 255, 1, 2, 3, 4, 5, 6, 7, 8},
+		{1, 2, 3, 4, 5, 6, 7, 8, 255, 1, 2, 3, 4, 5, 6, 7, 8, 255},
+		{1, 2, 3, 4, 5, 6, 7, 8, 255, 1, 2, 3, 4, 5, 6, 7, 8, 0},
+	}
+
+	for _, input := range errInputs {
+		_, _, err := decodeBytes(input)
+		require.Error(t, err)
+	}
+}
diff --git a/client/http/types.go b/client/http/types.go
@@ -299,6 +299,44 @@ func (r *Rule) Clone() *Rule {
 	return &clone
 }
 
+var (
+	_ json.Marshaler   = (*Rule)(nil)
+	_ json.Unmarshaler = (*Rule)(nil)
+)
+
+// MarshalJSON implements `json.Marshaler` interface to make sure we could set the correct start/end key.
+func (r *Rule) MarshalJSON() ([]byte, error) {
+	r.StartKeyHex = hex.EncodeToString(encodeBytes(r.StartKey))
+	r.EndKeyHex = hex.EncodeToString(encodeBytes(r.EndKey))
+	return json.Marshal(r)
+}
+
+// UnmarshalJSON implements `json.Unmarshaler` interface to make sure we could get the correct start/end key.
+func (r *Rule) UnmarshalJSON(bytes []byte) error {
+	if err := json.Unmarshal(bytes, r); err != nil {
+		return err
+	}
+
+	startKey, err := hex.DecodeString(r.StartKeyHex)
+	if err != nil {
+		return err
+	}
+
+	endKey, err := hex.DecodeString(r.EndKeyHex)
+	if err != nil {
+		return err
+	}
+
+	_, r.StartKey, err = decodeBytes(startKey)
+	if err != nil {
+		return err
+	}
+
+	_, r.EndKey, err = decodeBytes(endKey)
+
+	return err
+}
+
 // RuleOpType indicates the operation type
 type RuleOpType string
 
@@ -322,6 +360,44 @@ func (r RuleOp) String() string {
 	return string(b)
 }
 
+var (
+	_ json.Marshaler   = (*RuleOp)(nil)
+	_ json.Unmarshaler = (*RuleOp)(nil)
+)
+
+// MarshalJSON implements `json.Marshaler` interface to make sure we could set the correct start/end key.
+func (r *RuleOp) MarshalJSON() ([]byte, error) {
+	r.StartKeyHex = hex.EncodeToString(encodeBytes(r.StartKey))
+	r.EndKeyHex = hex.EncodeToString(encodeBytes(r.EndKey))
+	return json.Marshal(r)
+}
+
+// UnmarshalJSON implements `json.Unmarshaler` interface to make sure we could get the correct start/end key.
+func (r *RuleOp) UnmarshalJSON(bytes []byte) error {
+	if err := json.Unmarshal(bytes, r); err != nil {
+		return err
+	}
+
+	startKey, err := hex.DecodeString(r.StartKeyHex)
+	if err != nil {
+		return err
+	}
+
+	endKey, err := hex.DecodeString(r.EndKeyHex)
+	if err != nil {
+		return err
+	}
+
+	_, r.StartKey, err = decodeBytes(startKey)
+	if err != nil {
+		return err
+	}
+
+	_, r.EndKey, err = decodeBytes(endKey)
+
+	return err
+}
+
 // RuleGroup defines properties of a rule group.
 type RuleGroup struct {
 	ID       string `json:"id,omitempty"`