chore: optimize checksum calculation

Author: olavloite

.gitignore

@@ -1,2 +1,3 @@
 gorm/
 .idea
+.DS_Store

checksum_row_iterator.go

@@ -15,11 +15,14 @@
 package spannerdriver
 
 import (
-	"bytes"
 	"context"
 	"crypto/sha256"
+	"encoding/binary"
 	"encoding/gob"
+	"hash"
+	"math"
 	"reflect"
+	"sort"
 
 	"cloud.google.com/go/spanner"
 	sppb "cloud.google.com/go/spanner/apiv1/spannerpb"
@@ -70,8 +73,6 @@ type checksumRowIterator struct {
 	// seen. It is calculated as a SHA256 checksum over all rows that so far
 	// have been returned.
 	checksum *[32]byte
-	buffer   *bytes.Buffer
-	enc      *gob.Encoder
 
 	// errIndex and err indicate any error and the index in the result set
 	// where the error occurred.
@@ -110,7 +111,7 @@ func (it *checksumRowIterator) Next() (row *spanner.Row, err error) {
 			// checksum of the columns that are included in this result. This is
 			// also used to detect the possible difference between two empty
 			// result sets with a different set of columns.
-			it.checksum, err = createMetadataChecksum(it.enc, it.buffer, it.metadata)
+			it.checksum, err = createMetadataChecksum(it.metadata)
 			if err != nil {
 				return err
 			}
@@ -119,45 +120,87 @@ func (it *checksumRowIterator) Next() (row *spanner.Row, err error) {
 			return it.err
 		}
 		// Update the current checksum.
-		it.checksum, err = updateChecksum(it.enc, it.buffer, it.checksum, row)
+		it.checksum, err = updateChecksum(it.checksum, row)
 		return err
 	})
 	return row, err
 }
 
 // updateChecksum calculates the following checksum based on a current checksum
 // and a new row.
-func updateChecksum(enc *gob.Encoder, buffer *bytes.Buffer, currentChecksum *[32]byte, row *spanner.Row) (*[32]byte, error) {
-	buffer.Reset()
-	buffer.Write(currentChecksum[:])
+func updateChecksum(currentChecksum *[32]byte, row *spanner.Row) (*[32]byte, error) {
+	hasher := sha256.New()
+	hasher.Write(currentChecksum[:])
 	for i := 0; i < row.Size(); i++ {
 		var v spanner.GenericColumnValue
 		err := row.Column(i, &v)
 		if err != nil {
 			return nil, err
 		}
-		err = enc.Encode(v)
-		if err != nil {
-			return nil, err
+		hashValue(v.Value, hasher)
+	}
+	res := hasher.Sum(nil)
+	return (*[32]byte)(res), nil
+}
+
+var int32Buf [4]byte
+var float64Buf [8]byte
+
+func hashValue(value *structpb.Value, hasher hash.Hash) {
+	switch value.GetKind().(type) {
+	case *structpb.Value_StringValue:
+		hasher.Write(intToByte(int32Buf, len(value.GetStringValue())))
+		hasher.Write([]byte(value.GetStringValue()))
+	case *structpb.Value_NullValue:
+		hasher.Write([]byte{0})
+	case *structpb.Value_NumberValue:
+		hasher.Write(float64ToByte(float64Buf, value.GetNumberValue()))
+	case *structpb.Value_BoolValue:
+		if value.GetBoolValue() {
+			hasher.Write([]byte{1})
+		} else {
+			hasher.Write([]byte{0})
+		}
+	case *structpb.Value_StructValue:
+		fields := make([]string, 0, len(value.GetStructValue().Fields))
+		for field, _ := range value.GetStructValue().Fields {
+			fields = append(fields, field)
+		}
+		sort.Strings(fields)
+		for _, field := range fields {
+			hasher.Write(intToByte(int32Buf, len(field)))
+			hasher.Write([]byte(field))
+			hashValue(value.GetStructValue().Fields[field], hasher)
+		}
+	case *structpb.Value_ListValue:
+		for _, v := range value.GetListValue().GetValues() {
+			hashValue(v, hasher)
 		}
 	}
-	res := sha256.Sum256(buffer.Bytes())
-	return &res, nil
+}
+
+func intToByte(buf [4]byte, v int) []byte {
+	binary.BigEndian.PutUint32(buf[:], uint32(v))
+	return buf[:]
+}
+
+func float64ToByte(buf [8]byte, f float64) []byte {
+	binary.BigEndian.PutUint64(buf[:], math.Float64bits(f))
+	return buf[:]
 }
 
 // createMetadataChecksum calculates the checksum of the metadata of a result.
 // Only the column names and types are included in the checksum. Any transaction
 // metadata is not included.
-func createMetadataChecksum(enc *gob.Encoder, buffer *bytes.Buffer, metadata *sppb.ResultSetMetadata) (*[32]byte, error) {
-	buffer.Reset()
+func createMetadataChecksum(metadata *sppb.ResultSetMetadata) (*[32]byte, error) {
+	hasher := sha256.New()
 	for _, field := range metadata.RowType.Fields {
-		err := enc.Encode(field)
-		if err != nil {
-			return nil, err
-		}
+		hasher.Write(intToByte(int32Buf, len(field.Name)))
+		hasher.Write([]byte(field.Name))
+		hasher.Write(intToByte(int32Buf, int(field.Type.Code.Number())))
 	}
-	res := sha256.Sum256(buffer.Bytes())
-	return &res, nil
+	res := hasher.Sum(nil)
+	return (*[32]byte)(res), nil
 }
 
 // retry implements retriableStatement.retry for queries. It will execute the
@@ -167,8 +210,6 @@ func createMetadataChecksum(enc *gob.Encoder, buffer *bytes.Buffer, metadata *sp
 // initial iterator was also returned by the new iterator, and that the errors
 // were returned by the same row index.
 func (it *checksumRowIterator) retry(ctx context.Context, tx *spanner.ReadWriteStmtBasedTransaction) error {
-	buffer := &bytes.Buffer{}
-	enc := gob.NewEncoder(buffer)
 	retryIt := tx.QueryWithOptions(ctx, it.stmt, it.options)
 	// If the original iterator had been stopped, we should also always stop the
 	// new iterator.
@@ -198,7 +239,7 @@ func (it *checksumRowIterator) retry(ctx context.Context, tx *spanner.ReadWriteS
 	for n := int64(0); n < it.nc; n++ {
 		row, err := retryIt.Next()
 		if n == 0 && (err == nil || err == iterator.Done) {
-			newChecksum, checksumErr = createMetadataChecksum(enc, buffer, retryIt.Metadata)
+			newChecksum, checksumErr = createMetadataChecksum(retryIt.Metadata)
 			if checksumErr != nil {
 				return failRetry(checksumErr)
 			}
@@ -218,7 +259,7 @@ func (it *checksumRowIterator) retry(ctx context.Context, tx *spanner.ReadWriteS
 			}
 			return failRetry(ErrAbortedDueToConcurrentModification)
 		}
-		newChecksum, err = updateChecksum(enc, buffer, newChecksum, row)
+		newChecksum, err = updateChecksum(newChecksum, row)
 		if err != nil {
 			return failRetry(err)
 		}

checksum_row_iterator_test.go

@@ -15,8 +15,6 @@
 package spannerdriver
 
 import (
-	"bytes"
-	"encoding/gob"
 	"math/big"
 	"testing"
 	"time"
@@ -26,13 +24,6 @@ import (
 )
 
 func TestUpdateChecksum(t *testing.T) {
-	buffer1 := &bytes.Buffer{}
-	enc1 := gob.NewEncoder(buffer1)
-	buffer2 := &bytes.Buffer{}
-	enc2 := gob.NewEncoder(buffer2)
-	buffer3 := &bytes.Buffer{}
-	enc3 := gob.NewEncoder(buffer3)
-
 	row1, err := spanner.NewRow(
 		[]string{
 			"ColBool", "ColInt64", "ColFloat64", "ColNumeric", "ColString", "ColBytes", "ColDate", "ColTimestamp", "ColJson",
@@ -111,17 +102,17 @@ func TestUpdateChecksum(t *testing.T) {
 		t.Fatalf("could not create row 3: %v", err)
 	}
 	initial1 := new([32]byte)
-	checksum1, err := updateChecksum(enc1, buffer1, initial1, row1)
+	checksum1, err := updateChecksum(initial1, row1)
 	if err != nil {
 		t.Fatalf("could not calculate checksum 1: %v", err)
 	}
 	initial2 := new([32]byte)
-	checksum2, err := updateChecksum(enc2, buffer2, initial2, row2)
+	checksum2, err := updateChecksum(initial2, row2)
 	if err != nil {
 		t.Fatalf("could not calculate checksum 2: %v", err)
 	}
 	initial3 := new([32]byte)
-	checksum3, err := updateChecksum(enc3, buffer3, initial3, row3)
+	checksum3, err := updateChecksum(initial3, row3)
 	if err != nil {
 		t.Fatalf("could not calculate checksum 3: %v", err)
 	}
@@ -136,11 +127,11 @@ func TestUpdateChecksum(t *testing.T) {
 
 	// Updating checksums 1 and 3 with the data from row 2 should also produce
 	// the same checksum.
-	checksum1_2, err := updateChecksum(enc1, buffer1, checksum1, row2)
+	checksum1_2, err := updateChecksum(checksum1, row2)
 	if err != nil {
 		t.Fatalf("could not calculate checksum 1_2: %v", err)
 	}
-	checksum3_2, err := updateChecksum(enc3, buffer3, checksum3, row2)
+	checksum3_2, err := updateChecksum(checksum3, row2)
 	if err != nil {
 		t.Fatalf("could not calculate checksum 1_2: %v", err)
 	}
@@ -150,7 +141,7 @@ func TestUpdateChecksum(t *testing.T) {
 
 	// The combination of row 3 and 2 will produce a different checksum than the
 	// combination 2 and 3, because they are in a different order.
-	checksum2_3, err := updateChecksum(enc2, buffer2, checksum2, row3)
+	checksum2_3, err := updateChecksum(checksum2, row3)
 	if err != nil {
 		t.Fatalf("could not calculate checksum 2_3: %v", err)
 	}
@@ -160,9 +151,6 @@ func TestUpdateChecksum(t *testing.T) {
 }
 
 func TestUpdateChecksumForNullValues(t *testing.T) {
-	buffer := &bytes.Buffer{}
-	enc := gob.NewEncoder(buffer)
-
 	row, err := spanner.NewRow(
 		[]string{
 			"ColBool", "ColInt64", "ColFloat64", "ColNumeric", "ColString", "ColBytes", "ColDate", "ColTimestamp", "ColJson",
@@ -182,7 +170,7 @@ func TestUpdateChecksumForNullValues(t *testing.T) {
 	}
 	initial := new([32]byte)
 	// Create the initial checksum.
-	checksum, err := updateChecksum(enc, buffer, initial, row)
+	checksum, err := updateChecksum(initial, row)
 	if err != nil {
 		t.Fatalf("could not calculate checksum 1: %v", err)
 	}
@@ -192,14 +180,88 @@ func TestUpdateChecksumForNullValues(t *testing.T) {
 		t.Fatalf("checksum value should not be equal to the initial value")
 	}
 	// Calculating the same checksum again should yield the same result.
-	buffer2 := &bytes.Buffer{}
-	enc2 := gob.NewEncoder(buffer2)
 	initial2 := new([32]byte)
-	checksum2, err := updateChecksum(enc2, buffer2, initial2, row)
+	checksum2, err := updateChecksum(initial2, row)
 	if err != nil {
 		t.Fatalf("failed to update checksum: %v", err)
 	}
 	if *checksum != *checksum2 {
 		t.Fatalf("recalculated checksum does not match the initial calculation")
 	}
 }
+
+func BenchmarkChecksumRowIterator(b *testing.B) {
+	row1, _ := spanner.NewRow(
+		[]string{
+			"ColBool", "ColInt64", "ColFloat64", "ColNumeric", "ColString", "ColBytes", "ColDate", "ColTimestamp", "ColJson",
+			"ArrBool", "ArrInt64", "ArrFloat64", "ArrNumeric", "ArrString", "ArrBytes", "ArrDate", "ArrTimestamp", "ArrJson",
+		},
+		[]interface{}{
+			true, int64(1), 3.14, numeric("6.626"), "test", []byte("testbytes"), civil.Date{Year: 2021, Month: 8, Day: 5},
+			time.Date(2021, 8, 5, 13, 19, 23, 123456789, time.UTC),
+			nullJson(true, `"key": "value", "other-key": ["value1", "value2"]}`),
+			[]bool{true, false}, []int64{1, 2}, []float64{3.14, 6.626}, []big.Rat{numeric("3.14"), numeric("6.626")},
+			[]string{"test1", "test2"}, [][]byte{[]byte("testbytes1"), []byte("testbytes1")},
+			[]civil.Date{{Year: 2021, Month: 8, Day: 5}, {Year: 2021, Month: 8, Day: 6}},
+			[]time.Time{
+				time.Date(2021, 8, 5, 13, 19, 23, 123456789, time.UTC),
+				time.Date(2021, 8, 6, 13, 19, 23, 123456789, time.UTC),
+			},
+			[]spanner.NullJSON{
+				nullJson(true, `"key1": "value1", "other-key1": ["value1", "value2"]}`),
+				nullJson(true, `"key2": "value2", "other-key2": ["value1", "value2"]}`),
+			},
+		},
+	)
+	row2, _ := spanner.NewRow(
+		[]string{
+			"ColBool", "ColInt64", "ColFloat64", "ColNumeric", "ColString", "ColBytes", "ColDate", "ColTimestamp", "ColJson",
+			"ArrBool", "ArrInt64", "ArrFloat64", "ArrNumeric", "ArrString", "ArrBytes", "ArrDate", "ArrTimestamp", "ArrJson",
+		},
+		[]interface{}{
+			true, int64(2), 6.626, numeric("3.14"), "test2", []byte("testbytes2"), civil.Date{Year: 2020, Month: 8, Day: 5},
+			time.Date(2020, 8, 5, 13, 19, 23, 123456789, time.UTC),
+			nullJson(true, `"key": "other-value", "other-key": ["other-value1", "other-value2"]}`),
+			[]bool{true, false}, []int64{1, 2}, []float64{3.14, 6.626}, []big.Rat{numeric("3.14"), numeric("6.626")},
+			[]string{"test1_", "test2_"}, [][]byte{[]byte("testbytes1_"), []byte("testbytes1_")},
+			[]civil.Date{{Year: 2020, Month: 8, Day: 5}, {Year: 2020, Month: 8, Day: 6}},
+			[]time.Time{
+				time.Date(2020, 8, 5, 13, 19, 23, 123456789, time.UTC),
+				time.Date(2020, 8, 6, 13, 19, 23, 123456789, time.UTC),
+			},
+			[]spanner.NullJSON{
+				nullJson(true, `"key1": "other-value1", "other-key1": ["other-value1", "other-value2"]}`),
+				nullJson(true, `"key2": "other-value2", "other-key2": ["other-value1", "other-value2"]}`),
+			},
+		},
+	)
+	row3, _ := spanner.NewRow(
+		[]string{
+			"ColBool", "ColInt64", "ColFloat64", "ColNumeric", "ColString", "ColBytes", "ColDate", "ColTimestamp", "ColJson",
+			"ArrBool", "ArrInt64", "ArrFloat64", "ArrNumeric", "ArrString", "ArrBytes", "ArrDate", "ArrTimestamp", "ArrJson",
+		},
+		[]interface{}{
+			true, int64(1), 3.14, numeric("6.626"), "test", []byte("testbytes"), civil.Date{Year: 2021, Month: 8, Day: 5},
+			time.Date(2021, 8, 5, 13, 19, 23, 123456789, time.UTC),
+			nullJson(true, `"key": "value", "other-key": ["value1", "value2"]}`),
+			[]bool{true, false}, []int64{1, 2}, []float64{3.14, 6.626}, []big.Rat{numeric("3.14"), numeric("6.626")},
+			[]string{"test1", "test2"}, [][]byte{[]byte("testbytes1"), []byte("testbytes1")},
+			[]civil.Date{{Year: 2021, Month: 8, Day: 5}, {Year: 2021, Month: 8, Day: 6}},
+			[]time.Time{
+				time.Date(2021, 8, 5, 13, 19, 23, 123456789, time.UTC),
+				time.Date(2021, 8, 6, 13, 19, 23, 123456789, time.UTC),
+			},
+			[]spanner.NullJSON{
+				nullJson(true, `"key1": "value1", "other-key1": ["value1", "value2"]}`),
+				nullJson(true, `"key2": "value2", "other-key2": ["value1", "value2"]}`),
+			},
+		},
+	)
+
+	for b.Loop() {
+		initial := new([32]byte)
+		checksum, _ := updateChecksum(initial, row1)
+		checksum, _ = updateChecksum(checksum, row2)
+		checksum, _ = updateChecksum(checksum, row3)
+	}
+}