Prepend header bytes to WriteLog messages for backward compatibility (#2309)

MarshalLogMessage serializes the log message and prepends additional encoding
information to each message. Currently, we prepend three bytes to each log
message:
byte[0]: encodeVersion - this is currently a constant within the code. If we
modify structs such that they encode differently, we'll have to change the
constant and keep previous versions of structs for deserialization.
byte[1]: encodeType (e.g. "json", etc.)
byte[2]: logMessageType

If we get into a situation where we want more flexibility in these message
header bytes—for example, if we want to use more than three bytes—we could do
something with the first bit of the encodeVersion: if it's 1, that could
indicate that there are additional header bytes, and the following seven bits
could indicate how many.

(cherry picked from commit 6740bc250e)

dax/Makefile

@@ -15,11 +15,10 @@ testv:
 
 test-integration:
 	mkdir -p ../coverage-from-docker
-	$(GO) test ./test/dax -count 1 -run TestDAXIntegration
+	$(GO) test ./test/dax -count 1 -run TestDAXIntegration/$(RUN)
 
 testv-integration:
-	$(GO) test -v ./test/dax -count 1 -run TestDAXIntegration
-
+	$(GO) test -v ./test/dax -count 1 -run TestDAXIntegration/$(RUN)
 
 
 ############################### AWS STUFF ###############################

dax/computer/alpha/writelog.go

@@ -25,8 +25,15 @@ import (
 var _ computer.WriteLogReader = &alphaWriteLog{}
 var _ computer.WriteLogWriter = &alphaWriteLog{}
 
-// alphaWriteLog uses a WLer implementation (which could be, for example, an
-// http client or a locally running sub-service) to store its log messages.
+// alphaWriteLog is an implementation of the WriteLogReader and WriteLogWriter
+// interfaces. It uses a WriteLogger implementation (which could be, for
+// example, an http client or a locally running sub-service) to store its log
+// messages. I can't remember why we put this implementation in its own package
+// (a meaningless name called "alpha"); it probably has something to do with
+// wanting to adhere to a typical interface/implementation package structure,
+// but was confused by the current state of things where the "storage" package
+// (or perhaps "computer") is the top level package (i.e. "pilosa"). Until we
+// can correct the packaging, we will likely have weird cases like this.
 type alphaWriteLog struct {
 	wl featurebase.WriteLogger
 }
@@ -90,7 +97,7 @@ func (w *alphaWriteLog) DeleteFieldKeys(ctx context.Context, qtid dax.QualifiedT
 }
 
 func (w *alphaWriteLog) WriteShard(ctx context.Context, qtid dax.QualifiedTableID, partition dax.PartitionNum, shard dax.ShardNum, version int, msg computer.LogMessage) error {
-	b, err := computer.MarshalLogMessage(msg)
+	b, err := computer.MarshalLogMessage(msg, computer.EncodeTypeJSON)
 	if err != nil {
 		return errors.Wrap(err, "marshalling log message")
 	}
@@ -299,22 +306,7 @@ func (r *shardReader) Read() (computer.LogMessage, error) {
 	}
 
 	if r.scanner.Scan() {
-		b := r.scanner.Bytes()
-
-		if len(b) == 0 {
-			return nil, errors.New(errors.ErrUncoded, "empty log record")
-		}
-		logMessageType := b[0]
-
-		msg, err := computer.LogMessageByType(logMessageType)
-		if err != nil {
-			return nil, errors.Wrap(err, "getting log message by type")
-		}
-
-		if err := json.Unmarshal(b[1:], &msg); err != nil {
-			return nil, errors.Wrap(err, "unmarshaling log message")
-		}
-		return msg, nil
+		return computer.UnmarshalLogMessage(r.scanner.Bytes())
 	}
 	if err := r.scanner.Err(); err != nil {
 		return nil, err

dax/computer/writelog.go

@@ -17,14 +17,22 @@ type WriteLogWriter interface {
 	// CreateTableKeys sends a map of string key to uint64 ID for the table and
 	// partition provided.
 	CreateTableKeys(ctx context.Context, qtid dax.QualifiedTableID, partition dax.PartitionNum, version int, _ map[string]uint64) error
+
+	// DeleteTableKeys deletes all table keys for the table and partition
+	// provided.
 	DeleteTableKeys(ctx context.Context, qtid dax.QualifiedTableID, partition dax.PartitionNum, version int) error
 
 	// CreateFieldKeys sends a map of string key to uint64 ID for the table and
 	// field provided.
 	CreateFieldKeys(ctx context.Context, qtid dax.QualifiedTableID, field dax.FieldName, version int, _ map[string]uint64) error
+
+	// DeleteTableKeys deletes all field keys for the table and field provided.
 	DeleteFieldKeys(ctx context.Context, qtid dax.QualifiedTableID, field dax.FieldName, version int) error
 
+	// WriteShard sends shard data for the table and shard provided.
 	WriteShard(ctx context.Context, qtid dax.QualifiedTableID, partition dax.PartitionNum, shard dax.ShardNum, version int, msg LogMessage) error
+
+	// DeleteShard deletes all data for the table and shard provided.
 	DeleteShard(ctx context.Context, qtid dax.QualifiedTableID, partition dax.PartitionNum, shard dax.ShardNum, version int) error
 }
 
@@ -180,29 +188,111 @@ type FieldKeyMap struct {
 }
 
 const (
-	logMessageTypeImportRoaring = iota
+	logMessageTypeImportRoaring byte = iota
 	logMessageTypeImport
 	logMessageTypeImportValue
 	logMessageTypeImportRoaringShard
 )
 
+// encoderKey* are part of the log message header. They indicate what encoding
+// type a specific log message is serialized with.
+const (
+	encoderKeyJSON byte = iota
+)
+
+const (
+	EncodeTypeJSON string = "json"
+
+	// encodeVersion refers to the version of the structs used to represent the
+	// log messages. If we change structs, we'll need to modify this version
+	// number and maintain the previous version of the structs somewhere for
+	// deserialization.
+	encodeVersion byte = 1
+)
+
+// logMessageEncoder is implemented by any encoder used to serialize LogMessages
+// to []byte.
+type logMessageEncoder interface {
+	Key() byte
+	Marshal(LogMessage) ([]byte, error)
+	Unmarshal([]byte, LogMessage) error
+}
+
+// LogMessage is implemented by a variety of types which can be serialized as
+// messages to the WriteLogger.
 type LogMessage interface{}
 
-// MarshalLogMessage serializes the log message and adds log message type info.
-func MarshalLogMessage(msg LogMessage) ([]byte, error) {
-	typ, err := getLogMessageType(msg)
+// MarshalLogMessage serializes the log message and prepends additional encoding
+// information to each message. Currently, we prepend three bytes to each log
+// message:
+// byte[0]: encodeVersion - this is currently a constant within the code. If we
+// modify structs such that they encode differently, we'll have to change the
+// constant and keep previous versions of structs for deserialization.
+// byte[1]: encodeType (e.g. "json", etc.)
+// byte[2]: logMessageType
+//
+// If we get into a situation where we want more flexibility in these message
+// header bytes—for example, if we want to use more than three bytes—we could do
+// something with the first bit of the encodeVersion: if it's 1, that could
+// indicate that there are additional header bytes, and the following seven bits
+// could indicate how many.
+func MarshalLogMessage(msg LogMessage, encode string) ([]byte, error) {
+	encoder, err := getEncoderByType(encode)
+	if err != nil {
+		return nil, errors.Wrap(err, "getting encoder by type")
+	}
+
+	logMessageType, err := getLogMessageType(msg)
 	if err != nil {
 		return nil, errors.Wrap(err, "getting log message type")
 	}
 
-	buf, err := json.Marshal(msg)
+	var buf []byte
+
+	buf, err = encoder.Marshal(msg)
 	if err != nil {
 		return nil, errors.Wrap(err, "marshaling log message")
 	}
-	return append([]byte{typ}, buf...), nil
+
+	return append([]byte{encodeVersion, encoder.Key(), logMessageType}, buf...), nil
 }
 
-func LogMessageByType(typ byte) (LogMessage, error) {
+// UnmarshalLogMessage deserializes the log message based on the log message
+// type info.
+func UnmarshalLogMessage(b []byte) (LogMessage, error) {
+	if len(b) < 3 {
+		return nil, errors.New(errors.ErrUncoded, "log record does not contain a full header")
+	}
+
+	encVersion := b[0]
+	encKey := b[1]
+	logMessageType := b[2]
+
+	// Ensure that the log message is able to be handled by this code. If we
+	// increment the constant encodeVersion, we'll need to modify this to handle
+	// the log based on previous encodeVersions.
+	if encVersion != encodeVersion {
+		return nil, errors.Errorf("encode version is unsupported: %d", encVersion)
+	}
+
+	msg, err := logMessageByType(logMessageType)
+	if err != nil {
+		return nil, errors.Wrap(err, "getting log message by type")
+	}
+
+	encoder, err := getEncoderByKey(encKey)
+	if err != nil {
+		return nil, errors.Wrap(err, "getting encoder by key")
+	}
+
+	if err := encoder.Unmarshal(b[3:], &msg); err != nil {
+		return nil, errors.Wrap(err, "unmarshaling log message")
+	}
+
+	return msg, nil
+}
+
+func logMessageByType(typ byte) (LogMessage, error) {
 	switch typ {
 	case logMessageTypeImportRoaring:
 		return &ImportRoaringMessage{}, nil
@@ -232,6 +322,24 @@ func getLogMessageType(m LogMessage) (byte, error) {
 	}
 }
 
+func getEncoderByType(encode string) (logMessageEncoder, error) {
+	switch encode {
+	case EncodeTypeJSON:
+		return &encoderJSON{}, nil
+	default:
+		return nil, errors.Errorf("invalid encode type: %s", encode)
+	}
+}
+
+func getEncoderByKey(id byte) (logMessageEncoder, error) {
+	switch id {
+	case encoderKeyJSON:
+		return &encoderJSON{}, nil
+	default:
+		return nil, errors.Errorf("invalid encode type: %d", id)
+	}
+}
+
 type ImportRoaringMessage struct {
 	LogMessage `json:"-"`
 
@@ -305,3 +413,22 @@ type RoaringUpdate struct {
 	Set          []byte `json:"set"`
 	ClearRecords bool   `json:"clear-records"`
 }
+
+// Ensure type implements interface.
+var _ logMessageEncoder = (*encoderJSON)(nil)
+
+// encoderJSON is an implementation of the logMessageEncoder interface which
+// encodes LogMessages as JSON.
+type encoderJSON struct{}
+
+func (e *encoderJSON) Key() byte {
+	return encoderKeyJSON
+}
+
+func (e *encoderJSON) Marshal(msg LogMessage) ([]byte, error) {
+	return json.Marshal(msg)
+}
+
+func (e *encoderJSON) Unmarshal(b []byte, msg LogMessage) error {
+	return json.Unmarshal(b, msg)
+}

dax/computer/writelog_test.go

@@ -0,0 +1,154 @@
+package computer_test
+
+import (
+	"fmt"
+	"reflect"
+	"testing"
+
+	"github.com/molecula/featurebase/v3/dax/computer"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestComputer_WriteLog(t *testing.T) {
+	t.Run("Marshal", func(t *testing.T) {
+		tableName := "tbl"
+
+		type invalidType struct{}
+
+		tests := []struct {
+			msg            computer.LogMessage
+			encodeType     string
+			expVersion     byte
+			expEncodeType  byte
+			expMessageType byte
+			expError       string
+		}{
+			{
+				msg: &computer.ImportRoaringMessage{
+					Table: tableName,
+				},
+				encodeType:     computer.EncodeTypeJSON,
+				expVersion:     1,
+				expEncodeType:  0,
+				expMessageType: 0,
+			},
+			{
+				msg: &computer.ImportMessage{
+					Table: tableName,
+				},
+				encodeType:     computer.EncodeTypeJSON,
+				expVersion:     1,
+				expEncodeType:  0,
+				expMessageType: 1,
+			},
+			{
+				msg: &computer.ImportValueMessage{
+					Table: tableName,
+				},
+				encodeType:     computer.EncodeTypeJSON,
+				expVersion:     1,
+				expEncodeType:  0,
+				expMessageType: 2,
+			},
+			{
+				msg: &computer.ImportRoaringShardMessage{
+					Table: tableName,
+				},
+				encodeType:     computer.EncodeTypeJSON,
+				expVersion:     1,
+				expEncodeType:  0,
+				expMessageType: 3,
+			},
+
+			// Error cases.
+			{
+				msg:        &invalidType{},
+				encodeType: computer.EncodeTypeJSON,
+				expError:   "don't have type",
+			},
+			{
+				msg: &computer.ImportMessage{
+					Table: tableName,
+				},
+				encodeType: "badencoding",
+				expError:   "invalid encode type",
+			},
+		}
+		for i, test := range tests {
+			t.Run(fmt.Sprintf("test-%d", i), func(t *testing.T) {
+				b, err := computer.MarshalLogMessage(test.msg, test.encodeType)
+				if test.expError != "" {
+					if assert.Error(t, err) {
+						assert.Contains(t, err.Error(), test.expError)
+					}
+					return
+				} else {
+					assert.NoError(t, err)
+				}
+
+				// Version
+				assert.Equal(t, test.expVersion, b[0])
+
+				// EncodeType
+				assert.Equal(t, test.expEncodeType, b[1])
+
+				// MessageType
+				assert.Equal(t, test.expMessageType, b[2])
+
+				logMessage, err := computer.UnmarshalLogMessage(b)
+				assert.NoError(t, err)
+
+				// Get the value in Table for the instance of LogMessage.
+				val := reflect.ValueOf(logMessage).Elem()
+				fld := val.FieldByName("Table")
+
+				assert.Equal(t, tableName, fld.String())
+			})
+		}
+	})
+
+	t.Run("Unmarshal", func(t *testing.T) {
+		tests := []struct {
+			b        []byte
+			expError string
+		}{
+			{
+				b:        []byte{1, 0, 1, '{', '}'},
+				expError: "",
+			},
+			{
+				b:        []byte{},
+				expError: "log record does not contain a full header",
+			},
+			{
+				b:        []byte{1, 0, 1},
+				expError: "unexpected end of JSON input",
+			},
+			{
+				b:        []byte{1, 255, 1},
+				expError: "getting encoder by key: invalid encode type: 255",
+			},
+			{
+				b:        []byte{1, 0, 255},
+				expError: "unknown message type",
+			},
+			{
+				b:        []byte{255, 0, 1},
+				expError: "encode version is unsupported",
+			},
+		}
+		for i, test := range tests {
+			t.Run(fmt.Sprintf("test-%d", i), func(t *testing.T) {
+				_, err := computer.UnmarshalLogMessage(test.b)
+				if test.expError != "" {
+					if assert.Error(t, err) {
+						assert.Contains(t, err.Error(), test.expError)
+					}
+					return
+				} else {
+					assert.NoError(t, err)
+				}
+			})
+		}
+	})
+}