Use proto JSON serializer for proto objects (#611)

Author: alexshtin

common/archiver/historyIterator.go

@@ -25,9 +25,12 @@
 package archiver
 
 import (
+	"bytes"
 	"encoding/json"
 	"errors"
 
+	"github.com/gogo/protobuf/jsonpb"
+	"github.com/gogo/protobuf/proto"
 	historypb "go.temporal.io/api/history/v1"
 	"go.temporal.io/api/serviceerror"
 
@@ -236,19 +239,27 @@ type (
 		EstimateSize(v interface{}) (int, error)
 	}
 
-	jsonSizeEstimator struct{}
+	jsonSizeEstimator struct {
+		marshaler jsonpb.Marshaler
+	}
 )
 
 func (e *jsonSizeEstimator) EstimateSize(v interface{}) (int, error) {
+	// jsonpb must be used for proto structs.
+	if protoMessage, ok := v.(proto.Message); ok {
+		var buf bytes.Buffer
+		err := e.marshaler.Marshal(&buf, protoMessage)
+		return buf.Len(), err
+	}
+
 	data, err := json.Marshal(v)
 	if err != nil {
 		return 0, err
 	}
 	return len(data), nil
 }
 
-// NewJSONSizeEstimator returns a new SizeEstimator which uses json encoding to
-// estimate size
+// NewJSONSizeEstimator returns a new SizeEstimator which uses json encoding to estimate size
 func NewJSONSizeEstimator() SizeEstimator {
 	return &jsonSizeEstimator{}
 }

common/archiver/historyIterator_test.go

@@ -27,12 +27,15 @@ package archiver
 import (
 	"errors"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
 	"github.com/stretchr/testify/suite"
+	enumspb "go.temporal.io/api/enums/v1"
 	historypb "go.temporal.io/api/history/v1"
 	"go.temporal.io/api/serviceerror"
+	taskqueuepb "go.temporal.io/api/taskqueue/v1"
 
 	archiverspb "go.temporal.io/server/api/archiver/v1"
 	"go.temporal.io/server/common"
@@ -710,3 +713,32 @@ func (s *HistoryIteratorSuite) constructTestHistoryIterator(
 	itr.sizeEstimator = newTestSizeEstimator()
 	return itr
 }
+func (s *HistoryIteratorSuite) TestJSONSizeEstimator() {
+	e := NewJSONSizeEstimator()
+
+	historyEvent := &historypb.HistoryEvent{
+		EventId:   1,
+		Timestamp: time.Now().UnixNano(),
+		TaskId:    1,
+		Version:   1,
+	}
+	historyEvent.EventType = enumspb.EVENT_TYPE_WORKFLOW_TASK_SCHEDULED
+	historyEvent.Attributes = &historypb.HistoryEvent_WorkflowTaskScheduledEventAttributes{WorkflowTaskScheduledEventAttributes: &historypb.WorkflowTaskScheduledEventAttributes{
+		TaskQueue: &taskqueuepb.TaskQueue{
+			Name: "taskQueue",
+			Kind: enumspb.TASK_QUEUE_KIND_NORMAL,
+		},
+		StartToCloseTimeoutSeconds: 10,
+		Attempt:                    1,
+	}}
+
+	h := &historypb.History{
+		Events: []*historypb.HistoryEvent{
+			historyEvent,
+		},
+	}
+
+	size, err := e.EstimateSize(h)
+	s.NoError(err)
+	s.Equal(261, size)
+}

common/persistence/elasticsearch/decodeBench_test.go

@@ -48,16 +48,15 @@ var (
          "WorkflowId": "6bfbc1e5-6ce4-4e22-bbfb-e0faa9a7a604-1-2256",
          "WorkflowType": "TestWorkflowExecute",
  		 "Encoding" : "proto3",
-		  "TaskQueue" : "taskQueue", 
-		  "Memo" : "deadbeef====="}`)
+		 "TaskQueue" : "taskQueue", 
+		 "Memo" : "deadbeef====="}`)
 )
 
 /*
 BenchmarkJSONDecodeToType-8       200000              9321 ns/op
 BenchmarkJSONDecodeToMap-8        100000             12878 ns/op
 */
 
-//nolint
 func BenchmarkJSONDecodeToType(b *testing.B) {
 	bytes := (*json.RawMessage)(&data)
 	for i := 0; i < b.N; i++ {
@@ -71,14 +70,14 @@ func BenchmarkJSONDecodeToType(b *testing.B) {
 			ExecutionTime: time.Unix(0, source.ExecutionTime),
 			Memo:          p.NewDataBlob(source.Memo, common.EncodingType(source.Encoding)),
 			TaskQueue:     source.TaskQueue,
+			CloseTime:     time.Unix(0, source.CloseTime),
+			Status:        source.ExecutionStatus,
+			HistoryLength: source.HistoryLength,
 		}
-		record.CloseTime = time.Unix(0, source.CloseTime)
-		record.Status = source.ExecutionStatus
-		record.HistoryLength = source.HistoryLength
+		_ = record
 	}
 }
 
-//nolint
 func BenchmarkJSONDecodeToMap(b *testing.B) {
 	for i := 0; i < b.N; i++ {
 		var source map[string]interface{}

common/persistence/serializer.go

@@ -384,6 +384,7 @@ func (t *serializerImpl) serialize(input interface{}, encodingType common.Encodi
 		return nil, nil
 	}
 
+	// This should not pass proto struct down to JSON section.
 	if p, ok := input.(proto.Marshaler); ok {
 		return t.serializeProto(p, encodingType)
 	}
@@ -394,6 +395,7 @@ func (t *serializerImpl) serialize(input interface{}, encodingType common.Encodi
 	switch encodingType {
 	case common.EncodingTypeJSON, common.EncodingTypeUnknown, common.EncodingTypeEmpty: // For backward-compatibility
 		encodingType = common.EncodingTypeJSON
+		// input should never be a proto struct.
 		data, err = json.Marshal(input)
 	default:
 		return nil, NewUnknownEncodingTypeError(encodingType)
@@ -406,30 +408,6 @@ func (t *serializerImpl) serialize(input interface{}, encodingType common.Encodi
 	return NewDataBlob(data, encodingType), nil
 }
 
-func (t *serializerImpl) deserialize(data *serialization.DataBlob, target interface{}) error {
-	if data == nil {
-		return nil
-	}
-	if len(data.Data) == 0 {
-		return NewDeserializationError("DeserializeEvent empty data")
-	}
-	var err error
-
-	switch data.GetEncoding() {
-	case common.EncodingTypeProto3:
-		return NewDeserializationError(fmt.Sprintf("proto requires proto specific deserialization"))
-	case common.EncodingTypeJSON, common.EncodingTypeUnknown, common.EncodingTypeEmpty: // For backward-compatibility
-		err = json.Unmarshal(data.Data, target)
-	default:
-		return NewUnknownEncodingTypeError(data.GetEncoding())
-	}
-
-	if err != nil {
-		return NewDeserializationError(fmt.Sprintf("DeserializeBatchEvents encoding: \"%v\", error: %v", data.Encoding, err.Error()))
-	}
-	return nil
-}
-
 // NewUnknownEncodingTypeError returns a new instance of encoding type error
 func NewUnknownEncodingTypeError(encodingType common.EncodingType) error {
 	return &UnknownEncodingTypeError{encodingType: encodingType}

tools/cli/util.go

@@ -755,7 +755,7 @@ func processJSONInput(c *cli.Context) *commonpb.Payloads {
 	}
 	p, err := payloads.Encode(jsons...)
 	if err != nil {
-		ErrorAndExit("Unable to encode Input.", err)
+		ErrorAndExit("Unable to encode input.", err)
 	}
 
 	return p