br/lightning: add range properties codec and kv/stats reader

br/pkg/lightning/backend/external/codec.go

@@ -0,0 +1,78 @@
+// Copyright 2023 PingCAP, Inc.
+//
+// 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 external
+
+import "encoding/binary"
+
+type rangeProperty struct {
+	key    []byte
+	offset uint64
+	size   uint64
+	keys   uint64
+}
+
+// decodeMultiProps is only used for test.
+func decodeMultiProps(data []byte) ([]*rangeProperty, error) {
+	var ret []*rangeProperty
+	for len(data) > 0 {
+		propLen := int(binary.BigEndian.Uint32(data))
+		propBytes := data[4 : 4+propLen]
+		rp, err := decodeProp(propBytes)
+		if err != nil {
+			return nil, err
+		}
+		ret = append(ret, rp)
+		data = data[4+propLen:]
+	}
+	return ret, nil
+}
+
+func decodeProp(data []byte) (*rangeProperty, error) {
+	rp := &rangeProperty{}
+	keyLen := binary.BigEndian.Uint32(data[0:4])
+	rp.key = data[4 : 4+keyLen]
+	rp.size = binary.BigEndian.Uint64(data[4+keyLen : 12+keyLen])
+	rp.keys = binary.BigEndian.Uint64(data[12+keyLen : 20+keyLen])
+	rp.offset = binary.BigEndian.Uint64(data[20+keyLen : 28+keyLen])
+	return rp, nil
+}
+
+// keyLen + p.size + p.keys + p.offset
+const propertyLengthExceptKey = 4 + 8 + 8 + 8
+
+func encodeMultiProps(buf []byte, props []*rangeProperty) []byte {
+	var propLen [4]byte
+	for _, p := range props {
+		binary.BigEndian.PutUint32(propLen[:],
+			uint32(propertyLengthExceptKey+len(p.key)))
+		buf = append(buf, propLen[:4]...)
+		buf = encodeProp(buf, p)
+	}
+	return buf
+}
+
+func encodeProp(buf []byte, r *rangeProperty) []byte {
+	var b [8]byte
+	binary.BigEndian.PutUint32(b[:], uint32(len(r.key)))
+	buf = append(buf, b[:4]...)
+	buf = append(buf, r.key...)
+	binary.BigEndian.PutUint64(b[:], r.size)
+	buf = append(buf, b[:]...)
+	binary.BigEndian.PutUint64(b[:], r.keys)
+	buf = append(buf, b[:]...)
+	binary.BigEndian.PutUint64(b[:], r.offset)
+	buf = append(buf, b[:]...)
+	return buf
+}

br/pkg/lightning/backend/external/codec_test.go

@@ -0,0 +1,47 @@
+// Copyright 2023 PingCAP, Inc.
+//
+// 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 external
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestRangePropertyCodec(t *testing.T) {
+	prop := &rangeProperty{
+		key:    []byte("key"),
+		offset: 1,
+		size:   2,
+		keys:   3,
+	}
+	buf := encodeProp(nil, prop)
+	prop2, err := decodeProp(buf)
+	require.NoError(t, err)
+	require.EqualValues(t, prop, prop2)
+
+	p1, p2, p3 := &rangeProperty{}, &rangeProperty{}, &rangeProperty{}
+	for i, p := range []*rangeProperty{p1, p2, p3} {
+		p.key = []byte(fmt.Sprintf("key%d", i))
+		p.offset = uint64(10 * i)
+		p.size = uint64(20 * i)
+		p.keys = uint64(30 * i)
+	}
+	buf = encodeMultiProps(nil, []*rangeProperty{p1, p2, p3})
+	props, err := decodeMultiProps(buf)
+	require.NoError(t, err)
+	require.EqualValues(t, []*rangeProperty{p1, p2, p3}, props)
+}

br/pkg/lightning/backend/external/kv_reader.go

@@ -0,0 +1,67 @@
+// Copyright 2023 PingCAP, Inc.
+//
+// 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 external
+
+import (
+	"context"
+	"encoding/binary"
+
+	"github.com/pingcap/tidb/br/pkg/storage"
+)
+
+type kvReader struct {
+	byteReader *byteReader
+}
+
+func newKVReader(ctx context.Context, name string, store storage.ExternalStorage, initFileOffset uint64, bufSize int) (*kvReader, error) {
+	sr, err := openStoreReaderAndSeek(ctx, store, name, initFileOffset)
+	if err != nil {
+		return nil, err
+	}
+	br, err := newByteReader(ctx, sr, bufSize)
+	return &kvReader{
+		byteReader: br,
+	}, nil
+}
+
+func (r *kvReader) nextKV() (key, val []byte, err error) {
+	if r.byteReader.eof() {
+		return nil, nil, nil
+	}
+	r.byteReader.reset()
+	lenBytes, err := r.byteReader.readNBytes(8)
+	if err != nil {
+		return nil, nil, err
+	}
+	keyLen := int(binary.BigEndian.Uint64(*lenBytes))
+	keyPtr, err := r.byteReader.readNBytes(keyLen)
+	if err != nil {
+		return nil, nil, err
+	}
+	lenBytes, err = r.byteReader.readNBytes(8)
+	if err != nil {
+		return nil, nil, err
+	}
+	valLen := int(binary.BigEndian.Uint64(*lenBytes))
+	valPtr, err := r.byteReader.readNBytes(valLen)
+	if err != nil {
+		return nil, nil, err
+	}
+	return *keyPtr, *valPtr, nil
+}
+
+func (r *kvReader) Close() error {
+	return r.byteReader.Close()
+}

br/pkg/lightning/backend/external/stat_reader.go

@@ -0,0 +1,66 @@
+// Copyright 2023 PingCAP, Inc.
+//
+// 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 external
+
+import (
+	"context"
+	"encoding/binary"
+
+	"github.com/pingcap/tidb/br/pkg/storage"
+)
+
+type statsReader struct {
+	byteReader *byteReader
+	propBytes  []byte
+}
+
+func newStatsReader(ctx context.Context, store storage.ExternalStorage, name string, bufSize int) (*statsReader, error) {
+	sr, err := openStoreReaderAndSeek(ctx, store, name, 0)
+	if err != nil {
+		return nil, err
+	}
+	br, err := newByteReader(ctx, sr, bufSize)
+	if err != nil {
+		return nil, err
+	}
+	return &statsReader{
+		byteReader: br,
+		propBytes:  nil,
+	}, nil
+}
+
+func (r *statsReader) nextProp() (*rangeProperty, error) {
+	if r.byteReader.eof() {
+		return nil, nil
+	}
+	r.byteReader.reset()
+	lenBytes, err := r.byteReader.readNBytes(4)
+	if err != nil {
+		return nil, err
+	}
+	propLen := int(binary.BigEndian.Uint32(*lenBytes))
+	if cap(r.propBytes) < propLen {
+		r.propBytes = make([]byte, propLen)
+	}
+	propBytes, err := r.byteReader.readNBytes(propLen)
+	if err != nil {
+		return nil, err
+	}
+	return decodeProp(*propBytes)
+}
+
+func (r *statsReader) Close() error {
+	return r.byteReader.Close()
+}