feature: add basic tracing

Author: osmaczko

cmd/generate_handlers/generate_handlers_template.txt

@@ -4,7 +4,8 @@
 package protocol
 
 import (
-        "errors"
+	"context"
+	"errors"
 
 	"github.com/golang/protobuf/proto"
 	"go.uber.org/zap"
@@ -15,20 +16,24 @@ import (
 	messagingtypes "github.com/status-im/status-go/messaging/types"
 )
 
-func (m *Messenger) dispatchToHandler(messageState *ReceivedMessageState, protoBytes []byte, msg *common.StatusMessage, filter messagingtypes.ChatFilter, fromArchive bool) error {
+func (m *Messenger) dispatchToHandler(ctx context.Context, messageState *ReceivedMessageState, protoBytes []byte, msg *common.StatusMessage, filter messagingtypes.ChatFilter, fromArchive bool) error {
 	switch msg.ApplicationLayer.Type {
 	{{ range .}}
            case protobuf.ApplicationMetadataMessage_{{.EnumValue}}:
-		return m.{{.MethodName}}(messageState, protoBytes, msg, filter{{ if .FromArchiveArg }}, fromArchive{{ end }})
+		return m.{{.MethodName}}(ctx, messageState, protoBytes, msg, filter{{ if .FromArchiveArg }}, fromArchive{{ end }})
         {{ end }}
 	}
 	m.logger.Info("protobuf type not found", zap.String("type", string(msg.ApplicationLayer.Type)))
 	return errors.New("protobuf type not found")
 }
 
 {{ range . }}
-func (m *Messenger) {{.MethodName}}(messageState *ReceivedMessageState, protoBytes []byte, msg *common.StatusMessage, filter messagingtypes.ChatFilter{{ if .FromArchiveArg }}, fromArchive bool{{ end }}) error {
+func (m *Messenger) {{.MethodName}}(ctx context.Context, messageState *ReceivedMessageState, protoBytes []byte, msg *common.StatusMessage, filter messagingtypes.ChatFilter{{ if .FromArchiveArg }}, fromArchive bool{{ end }}) error {
 	m.logger.Info("handling {{ .ProtobufName}}")
+
+	ctx, span := m.tracer.Start(ctx, "Messenger.handle{{.ProtobufName}}")
+	defer span.End()
+
 	{{ if .SyncMessage }}
 	if !crypto.IsPubKeyEqual(messageState.CurrentMessageState.PublicKey, &m.identity.PublicKey) {
 		m.logger.Warn("not coming from us, ignoring")
@@ -37,7 +42,7 @@ func (m *Messenger) {{.MethodName}}(messageState *ReceivedMessageState, protoByt
 	{{ end }}
 
 	{{ if .ProcessRaw }}
-	return m.Handle{{.ProtobufName}}(messageState, protoBytes, msg)
+	return m.Handle{{.ProtobufName}}(ctx, messageState, protoBytes, msg)
 	{{ else }}
 	p := &protobuf.{{.ProtobufName}}{}
 	err := proto.Unmarshal(protoBytes, p)
@@ -47,7 +52,7 @@ func (m *Messenger) {{.MethodName}}(messageState *ReceivedMessageState, protoByt
 
 	m.outputToCSV(msg.TransportLayer.Message.Timestamp, msg.ApplicationLayer.ID, messageState.CurrentMessageState.Contact.ID, filter.ContentTopic(), filter.ChatID(), msg.ApplicationLayer.Type, p)
 
-	return m.Handle{{.ProtobufName}}(messageState, p, msg{{ if .FromArchiveArg }}, fromArchive {{ end }})
+	return m.Handle{{.ProtobufName}}(ctx, messageState, p, msg{{ if .FromArchiveArg }}, fromArchive {{ end }})
 	{{ end }}
 }
 

messaging/api.go

@@ -357,13 +357,13 @@ func (a *API) DecryptCommunityGrant(myIdentityKey *ecdsa.PrivateKey, senderKey *
 	return a.core.stack.Encryption.DecryptCommunityGrant(myIdentityKey, senderKey, grants)
 }
 
-func (a *API) HandleHashRatchetKeysPayload(groupID, encodedKeys []byte, myIdentityKey *ecdsa.PrivateKey, theirIdentityKey *ecdsa.PublicKey) error {
-	_, err := a.core.stack.Encryption.HandleHashRatchetKeysPayload(groupID, encodedKeys, myIdentityKey, theirIdentityKey)
+func (a *API) HandleHashRatchetKeysPayload(ctx context.Context, groupID, encodedKeys []byte, myIdentityKey *ecdsa.PrivateKey, theirIdentityKey *ecdsa.PublicKey) error {
+	_, err := a.core.stack.Encryption.HandleHashRatchetKeysPayload(ctx, groupID, encodedKeys, myIdentityKey, theirIdentityKey)
 	return err
 }
 
-func (a *API) HandleHashRatchetHeadersPayload(encodedHeaders [][]byte) error {
-	return a.core.stack.Encryption.HandleHashRatchetHeadersPayload(encodedHeaders)
+func (a *API) HandleHashRatchetHeadersPayload(ctx context.Context, encodedHeaders [][]byte) error {
+	return a.core.stack.Encryption.HandleHashRatchetHeadersPayload(ctx, encodedHeaders)
 }
 
 func (a *API) AddInstallation(identity []byte, timestamp int64, installation *types.Installation, enabled bool) ([]*types.Installation, error) {

messaging/controller/controller.go

@@ -11,6 +11,7 @@ import (
 
 	gocommon "github.com/status-im/status-go/common"
 	"github.com/status-im/status-go/crypto"
+	"github.com/status-im/status-go/internal/instrumentation/trace"
 	"github.com/status-im/status-go/messaging/adapters"
 	"github.com/status-im/status-go/messaging/common"
 	"github.com/status-im/status-go/messaging/controller/processor"
@@ -43,12 +44,13 @@ func NewController(
 	hashRatchetStorage common.HashRatchetPersistence,
 	publisher *pubsub.Publisher,
 	logger *zap.Logger,
+	tracer trace.Tracer,
 ) *Controller {
 	return &Controller{
 		identity:                   identity,
 		stack:                      stack,
-		sender:                     sender.NewSender(identity, stack, logger),
-		processor:                  processor.NewProcessor(identity, stack, messageConfirmationStorage, hashRatchetStorage, logger),
+		sender:                     sender.NewSender(identity, stack, logger, tracer),
+		processor:                  processor.NewProcessor(identity, stack, messageConfirmationStorage, hashRatchetStorage, logger, tracer),
 		messageConfirmationStorage: messageConfirmationStorage,
 		hashRatchetStorage:         hashRatchetStorage,
 		publisher:                  publisher,

messaging/controller/processor/processor.go

@@ -1,18 +1,23 @@
 package processor
 
 import (
+	"context"
 	"crypto/ecdsa"
 	"encoding/hex"
 
 	"github.com/golang/protobuf/proto"
 	"github.com/pkg/errors"
+	otelattribute "go.opentelemetry.io/otel/attribute"
+	oteltrace "go.opentelemetry.io/otel/trace"
 	"go.uber.org/zap"
 
 	"github.com/status-im/status-go/crypto"
 	cryptotypes "github.com/status-im/status-go/crypto/types"
 	ethtypes "github.com/status-im/status-go/eth-node/types"
+	"github.com/status-im/status-go/internal/instrumentation/trace"
 	"github.com/status-im/status-go/messaging/adapters"
 	"github.com/status-im/status-go/messaging/common"
+	"github.com/status-im/status-go/messaging/controller/utils"
 	"github.com/status-im/status-go/messaging/layers/encryption"
 	"github.com/status-im/status-go/messaging/layers/encryption/sharedsecret"
 	"github.com/status-im/status-go/messaging/layers/segmentation"
@@ -37,6 +42,8 @@ type Processor struct {
 
 	publisher *pubsub.Publisher
 	logger    *zap.Logger
+
+	tracer trace.Tracer
 }
 
 func NewProcessor(
@@ -45,6 +52,7 @@ func NewProcessor(
 	messageConfirmationStorage common.MessageConfirmationPersistence,
 	hashRatchetStorage common.HashRatchetPersistence,
 	logger *zap.Logger,
+	tracer trace.Tracer,
 ) *Processor {
 	return &Processor{
 		identity:                   identity,
@@ -54,6 +62,7 @@ func NewProcessor(
 		ephemeralKeysManager:       NewEphemeralKeysManager(maxNumOfEphemeralKeys),
 		publisher:                  pubsub.NewPublisher(),
 		logger:                     logger.Named("processor"),
+		tracer:                     tracer,
 	}
 }
 
@@ -122,29 +131,49 @@ func (r *Processor) processMessage(m *types.ReceivedMessage) (*processMessageRes
 		return nil, err
 	}
 
+	hashes := [][]byte{m.Hash}
 	if responseMessage.SegmentationLayer.Segmented {
 		// Segments not completed yet, stop processing
 		if !responseMessage.SegmentationLayer.Completed {
 			return nil, nil
 		}
+		hashes = responseMessage.SegmentationLayer.Hashes
 	}
 
-	err = r.processEncryptionLayer(responseMessage, logger)
-	if err != nil {
-		logger.Debug("failed to process encryption layer", zap.Error(err))
+	ctx, span := r.tracer.Start(trace.DeriveRemoteContext(utils.MergeByteSlices(hashes)), "Processor.processMessage",
+		oteltrace.WithAttributes(
+			otelattribute.String("hash", cryptotypes.EncodeHex(m.Hash)),
+			otelattribute.StringSlice("hashes", cryptotypes.EncodeHexes(hashes)),
+		),
+	)
+	defer span.End()
 
+	err = r.processEncryptionLayer(ctx, responseMessage, logger)
+	if err == nil {
+		span.AddEvent("encryption layer processed")
+	} else {
 		// Hash ratchet with a group id not found yet, save the message for future processing
 		if err == encryption.ErrHashRatchetGroupIDNotFound && len(responseMessage.EncryptionLayer.HashRatchetInfo) == 1 {
 			info := responseMessage.EncryptionLayer.HashRatchetInfo[0]
+			span.AddEvent("hash ratchet with group id not found yet", oteltrace.WithAttributes(
+				otelattribute.String("groupID", cryptotypes.ToHex(info.GroupID)),
+			))
 			return nil, r.hashRatchetStorage.SaveMessage(info.GroupID, info.KeyID, m)
+		} else {
+			span.AddEvent("encryption layer not processed", oteltrace.WithAttributes(
+				otelattribute.String("error", err.Error()),
+			))
+			logger.Debug("failed to process encryption layer", zap.Error(err))
 		}
 	}
 
 	messages, ackedMessageIDs, err := r.processReliabilityLayer(responseMessage, logger)
 	if err == nil {
+		span.AddEvent("reliability layer processed")
 		response.messages = messages
 		response.ackedMessageIDs = ackedMessageIDs
 	} else {
+		span.AddEvent("reliability layer not processed")
 		logger.Debug("failed to process reliability layer", zap.Error(err))
 	}
 
@@ -211,38 +240,45 @@ func processTransportLayer(m *types.Message, receivedMessage *types.ReceivedMess
 	return nil
 }
 
-func (r *Processor) processSegmentationLayer(m *types.Message) (segmented, completed bool, err error) {
-	var reconstructedPayload []byte
-	reconstructedPayload, err = r.stack.Segmentation.Reconstruct(m.TransportLayer.Payload, m.TransportLayer.SigPubKey)
+func (r *Processor) processSegmentationLayer(m *types.Message) error {
+	reconstructedPayload, transportIDs, err := r.stack.Segmentation.Reconstruct(
+		m.TransportLayer.Payload,
+		m.TransportLayer.SigPubKey,
+		m.TransportLayer.Hash)
 
 	switch err {
 	case nil:
 		m.TransportLayer.Payload = reconstructedPayload
-		segmented = true
-		completed = true
+		m.SegmentationLayer.Segmented = true
+		m.SegmentationLayer.Completed = true
+		m.SegmentationLayer.Hashes = transportIDs
 	case segmentation.ErrIncomplete:
-		segmented = true
-		completed = false
+		m.SegmentationLayer.Segmented = true
+		m.SegmentationLayer.Completed = false
 		err = nil
 	case segmentation.ErrInvalidPayload:
-		segmented = false
-		completed = false
+		m.SegmentationLayer.Segmented = false
+		m.SegmentationLayer.Completed = false
 		err = nil
 	}
 
-	return
+	return err
 }
 
-func (r *Processor) processEncryptionLayer(m *types.Message, logger *zap.Logger) error {
+func (r *Processor) processEncryptionLayer(ctx context.Context, m *types.Message, logger *zap.Logger) error {
 	logger = logger.Named("processEncryptionLayer")
 
+	ctx, span := r.tracer.Start(ctx, "Processor.processEncryptionLayer")
+	defer span.End()
+
 	// As we handle non-encrypted messages, we make sure that DecryptPayload
 	// is set regardless of whether this step is successful
 	m.EncryptionLayer.Payload = m.TransportLayer.Payload
 
 	// if it's an ephemeral key, we don't negotiate a topic
 	ephemeralKey := r.ephemeralKeysManager.GetPrivateKeyFor(m.TransportLayer.Dst)
 	if ephemeralKey != nil {
+		span.AddEvent("targeted ephemeral key")
 		return nil
 	}
 
@@ -253,6 +289,7 @@ func (r *Processor) processEncryptionLayer(m *types.Message, logger *zap.Logger)
 	}
 
 	response, err := r.stack.Encryption.HandleMessage(
+		ctx,
 		r.identity,
 		m.SigPubKey(),
 		&protocolMessage,

messaging/controller/processor/processor_test.go

@@ -1,6 +1,7 @@
 package processor
 
 import (
+	"context"
 	"crypto/ecdsa"
 	"math"
 	"testing"
@@ -16,6 +17,7 @@ import (
 	"go.uber.org/zap"
 
 	"github.com/status-im/status-go/crypto"
+	"github.com/status-im/status-go/internal/instrumentation/trace"
 	"github.com/status-im/status-go/messaging/common"
 	commonmigrations "github.com/status-im/status-go/messaging/common/migrations"
 	"github.com/status-im/status-go/messaging/layers/encryption"
@@ -109,6 +111,7 @@ func (s *ProcessorSuite) SetupTest() {
 		encryption.NewSQLitePersistence(db),
 		"installation-1",
 		s.logger,
+		trace.NewNoopTracer(),
 	)
 
 	stack.Reliability = reliability.NewReliability(
@@ -126,6 +129,7 @@ func (s *ProcessorSuite) SetupTest() {
 		common.NewSQLiteMessageConfirmationPersistence(db),
 		common.NewSQLiteHashRatchetPersistence(db),
 		s.logger,
+		trace.NewNoopTracer(),
 	)
 }
 
@@ -198,6 +202,7 @@ func (s *ProcessorSuite) TestProcessMessageDatasyncEncrypted() {
 		encryption.NewSQLitePersistence(senderDatabase),
 		"installation-2",
 		s.logger,
+		trace.NewNoopTracer(),
 	)
 
 	messageSpec, err := senderEncryptionProtocol.BuildEncryptedMessage(
@@ -243,14 +248,15 @@ func (s *ProcessorSuite) TestHandleOutOfOrderHashRatchet() {
 		encryption.NewSQLitePersistence(senderDatabase),
 		"installation-2",
 		s.logger,
+		trace.NewNoopTracer(),
 	)
 
 	ratchet, err := senderEncryptionProtocol.GenerateHashRatchetKey(groupID)
 	s.Require().NoError(err)
 
 	ratchets := []*encryption.HashRatchetKeyCompatibility{ratchet}
 
-	hashRatchetKeyExchangeMessage, err := senderEncryptionProtocol.BuildHashRatchetKeyExchangeMessage(senderKey, &s.processor.identity.PublicKey, groupID, ratchets)
+	hashRatchetKeyExchangeMessage, err := senderEncryptionProtocol.BuildHashRatchetKeyExchangeMessage(context.Background(), senderKey, &s.processor.identity.PublicKey, groupID, ratchets)
 	s.Require().NoError(err)
 
 	encryptedPayload1, err := proto.Marshal(hashRatchetKeyExchangeMessage.Message)

messaging/controller/sender/sender.go

@@ -1,13 +1,21 @@
 package sender
 
 import (
+	"context"
 	"crypto/ecdsa"
+	"runtime"
+	"runtime/debug"
 
 	"github.com/golang/protobuf/proto"
+	otelattribute "go.opentelemetry.io/otel/attribute"
+	oteltrace "go.opentelemetry.io/otel/trace"
 	"go.uber.org/zap"
 
+	cryptotypes "github.com/status-im/status-go/crypto/types"
 	ethtypes "github.com/status-im/status-go/eth-node/types"
+	"github.com/status-im/status-go/internal/instrumentation/trace"
 	"github.com/status-im/status-go/messaging/common"
+	"github.com/status-im/status-go/messaging/controller/utils"
 	"github.com/status-im/status-go/messaging/layers/encryption"
 	"github.com/status-im/status-go/pkg/pubsub"
 )
@@ -18,17 +26,20 @@ type Sender struct {
 
 	publisher *pubsub.Publisher
 	logger    *zap.Logger
+	tracer    trace.Tracer
 }
 
 func NewSender(
 	identity *ecdsa.PrivateKey,
 	stack *common.MessagingStack,
 	logger *zap.Logger,
+	tracer trace.Tracer,
 ) *Sender {
 	return &Sender{
 		identity:  identity,
 		stack:     stack,
 		logger:    logger.Named("sender"),
+		tracer:    tracer,
 		publisher: pubsub.NewPublisher(),
 	}
 }
@@ -62,3 +73,34 @@ func (s *Sender) processAndMarshalMessageSpec(spec *encryption.ProtocolMessageSp
 
 	return messageBytes, sharedSecretKey, nil
 }
+
+func (s *Sender) telemetryTraceMessage(ctx context.Context, hashes [][]byte) context.Context {
+	if !s.tracer.Enabled() {
+		return ctx
+	}
+
+	funcName := "Sender.unknown"
+	pc, _, _, ok := runtime.Caller(1)
+	if ok {
+		fn := runtime.FuncForPC(pc)
+		if fn != nil {
+			funcName = fn.Name()
+		}
+	}
+
+	linkSpanCtx := trace.DeriveSpanContext(utils.MergeByteSlices(hashes), false)
+
+	spanOpts := []oteltrace.SpanStartOption{
+		oteltrace.WithAttributes(
+			otelattribute.String("traceID", linkSpanCtx.TraceID().String()),
+			otelattribute.StringSlice("hashes", cryptotypes.EncodeHexes(hashes)),
+			otelattribute.String("stack", string(debug.Stack())),
+		),
+		oteltrace.WithLinks(oteltrace.Link{SpanContext: linkSpanCtx}),
+	}
+
+	newCtx, span := s.tracer.Start(ctx, funcName, spanOpts...)
+	span.End()
+
+	return newCtx
+}