Add command monitoring to core/connection

GODRIVER-19

Change-Id: I0742c8d4bcc7da76c5cd146ca8093e2a351554f8

Author: Divjot Arora

core/connection/connection.go

@@ -24,13 +24,16 @@ import (
 	"sync/atomic"
 	"time"
 
+	"github.com/mongodb/mongo-go-driver/bson"
 	"github.com/mongodb/mongo-go-driver/core/address"
 	"github.com/mongodb/mongo-go-driver/core/compressor"
 	"github.com/mongodb/mongo-go-driver/core/description"
+	"github.com/mongodb/mongo-go-driver/core/event"
 	"github.com/mongodb/mongo-go-driver/core/wiremessage"
 )
 
 var globalClientConnectionID uint64
+var emptyDoc = bson.NewDocument()
 
 func nextClientConnectionID() uint64 {
 	return atomic.AddUint64(&globalClientConnectionID, 1)
@@ -89,10 +92,12 @@ type connection struct {
 	compressor  compressor.Compressor // use for compressing messages
 	// server can compress response with any compressor supported by driver
 	compressorMap    map[wiremessage.CompressorID]compressor.Compressor
+	commandMap       map[int64]*event.CommandMetadata // map for monitoring commands sent to server
 	dead             bool
 	idleTimeout      time.Duration
 	idleDeadline     time.Time
 	lifetimeDeadline time.Time
+	cmdMonitor       *event.CommandMonitor
 	readTimeout      time.Duration
 	uncompressBuf    []byte // buffer to uncompress messages
 	writeTimeout     time.Duration
@@ -140,6 +145,7 @@ func New(ctx context.Context, addr address.Address, opts ...Option) (Connection,
 		conn:             nc,
 		compressBuf:      make([]byte, 256),
 		compressorMap:    compressorMap,
+		commandMap:       make(map[int64]*event.CommandMetadata),
 		addr:             addr,
 		idleTimeout:      cfg.idleTimeout,
 		lifetimeDeadline: lifetimeDeadline,
@@ -180,6 +186,7 @@ func New(ctx context.Context, addr address.Address, opts ...Option) (Connection,
 		desc = &d
 	}
 
+	c.cmdMonitor = cfg.cmdMonitor // attach the command monitor later to avoid monitoring auth
 	return c, desc, nil
 }
 
@@ -342,6 +349,223 @@ func (c *connection) uncompressMessage(compressed wiremessage.Compressed) ([]byt
 	return fullMessage, origHeader.OpCode, nil
 }
 
+func canMonitor(cmd string) bool {
+	if cmd == "authenticate" || cmd == "saslStart" || cmd == "saslContinue" || cmd == "getnonce" || cmd == "createUser" ||
+		cmd == "updateUser" || cmd == "copydbgetnonce" || cmd == "copydbsaslstart" || cmd == "copydb" {
+		return false
+	}
+
+	return true
+}
+
+func (c *connection) commandStartedEvent(wm wiremessage.WireMessage) error {
+	if c.cmdMonitor == nil || c.cmdMonitor.Started == nil {
+		return nil
+	}
+
+	startedEvent := &event.CommandStartedEvent{
+		ConnectionID: c.id,
+	}
+
+	var cmd *bson.Document
+	var err error
+
+	var acknowledged bool
+	switch converted := wm.(type) {
+	case wiremessage.Query:
+		cmd, err = bson.ReadDocument([]byte(converted.Query))
+		if err != nil {
+			return err
+		}
+
+		acknowledged = converted.AcknowledgedWrite()
+		startedEvent.DatabaseName = converted.FullCollectionName[:len(converted.FullCollectionName)-5] // remove $.cmd
+		startedEvent.RequestID = int64(converted.MsgHeader.RequestID)
+
+		cmdElem := cmd.ElementAt(0)
+		if cmdElem.Key() == "$query" {
+			cmd = cmdElem.Value().MutableDocument()
+		}
+	case wiremessage.Msg:
+		cmd, err = converted.GetMainDocument()
+		if err != nil {
+			return err
+		}
+
+		acknowledged = converted.AcknowledgedWrite()
+		arr, identifier, err := converted.GetSequenceArray()
+		if err != nil {
+			return err
+		}
+		if arr != nil {
+			cmd = cmd.Copy() // make copy to avoid changing original command
+			cmd.Append(bson.EC.Array(identifier, arr))
+		}
+
+		dbVal, err := cmd.LookupErr("$db")
+		if err != nil {
+			return err
+		}
+
+		startedEvent.DatabaseName = dbVal.StringValue()
+		startedEvent.RequestID = int64(converted.MsgHeader.RequestID)
+	}
+
+	startedEvent.Command = cmd
+	startedEvent.CommandName = cmd.ElementAt(0).Key()
+	if !canMonitor(startedEvent.CommandName) {
+		startedEvent.Command = emptyDoc
+	}
+
+	c.cmdMonitor.Started(startedEvent)
+
+	if !acknowledged {
+		if c.cmdMonitor.Succeeded == nil {
+			return nil
+		}
+
+		// unack writes must provide a CommandSucceededEvent with an { ok: 1 } reply
+		finishedEvent := event.CommandFinishedEvent{
+			DurationNanos: 0,
+			CommandName:   startedEvent.CommandName,
+			RequestID:     startedEvent.RequestID,
+			ConnectionID:  c.id,
+		}
+
+		c.cmdMonitor.Succeeded(&event.CommandSucceededEvent{
+			CommandFinishedEvent: finishedEvent,
+			Reply: bson.NewDocument(
+				bson.EC.Int32("ok", 1),
+			),
+		})
+
+		return nil
+	}
+
+	c.commandMap[startedEvent.RequestID] = event.CreateMetadata(startedEvent.CommandName)
+	return nil
+}
+
+func processReply(reply *bson.Document) (bool, string) {
+	iter := reply.Iterator()
+	var success bool
+	var errmsg string
+	var errCode int32
+
+	for iter.Next() {
+		elem := iter.Element()
+		switch elem.Key() {
+		case "ok":
+			switch elem.Value().Type() {
+			case bson.TypeInt32:
+				if elem.Value().Int32() == 1 {
+					success = true
+				}
+			case bson.TypeInt64:
+				if elem.Value().Int64() == 1 {
+					success = true
+				}
+			case bson.TypeDouble:
+				if elem.Value().Double() == 1 {
+					success = true
+				}
+			}
+		case "errmsg":
+			if str, ok := elem.Value().StringValueOK(); ok {
+				errmsg = str
+			}
+		case "code":
+			if c, ok := elem.Value().Int32OK(); ok {
+				errCode = c
+			}
+		}
+	}
+
+	if success {
+		return true, ""
+	}
+
+	fullErrMsg := fmt.Sprintf("Error code %d: %s", errCode, errmsg)
+	return false, fullErrMsg
+}
+
+func (c *connection) commandFinishedEvent(wm wiremessage.WireMessage) error {
+	if c.cmdMonitor == nil {
+		return nil
+	}
+
+	var reply *bson.Document
+	var requestID int64
+	var err error
+
+	switch converted := wm.(type) {
+	case wiremessage.Reply:
+		requestID = int64(converted.MsgHeader.ResponseTo)
+		reply, err = converted.GetMainDocument()
+	case wiremessage.Msg:
+		requestID = int64(converted.MsgHeader.ResponseTo)
+		reply, err = converted.GetMainDocument()
+	}
+
+	if err != nil {
+		return err
+	}
+
+	cmdMetadata := c.commandMap[requestID]
+	delete(c.commandMap, requestID)
+	success, errmsg := processReply(reply)
+
+	if (success && c.cmdMonitor.Succeeded == nil) || (!success && c.cmdMonitor.Failed == nil) {
+		return nil
+	}
+
+	finishedEvent := event.CommandFinishedEvent{
+		DurationNanos: cmdMetadata.TimeDifference(),
+		CommandName:   cmdMetadata.Name,
+		RequestID:     requestID,
+		ConnectionID:  c.id,
+	}
+
+	if success {
+		if !canMonitor(finishedEvent.CommandName) {
+			successEvent := &event.CommandSucceededEvent{
+				Reply:                emptyDoc,
+				CommandFinishedEvent: finishedEvent,
+			}
+			c.cmdMonitor.Succeeded(successEvent)
+			return nil
+		}
+
+		// if response has type 1 document sequence, the sequence must be included as a BSON array in the event's reply.
+		if opmsg, ok := wm.(wiremessage.Msg); ok {
+			arr, identifier, err := opmsg.GetSequenceArray()
+			if err != nil {
+				return err
+			}
+			if arr != nil {
+				reply = reply.Copy() // make copy to avoid changing original command
+				reply.Append(bson.EC.Array(identifier, arr))
+			}
+		}
+
+		successEvent := &event.CommandSucceededEvent{
+			Reply:                reply,
+			CommandFinishedEvent: finishedEvent,
+		}
+
+		c.cmdMonitor.Succeeded(successEvent)
+		return nil
+	}
+
+	failureEvent := &event.CommandFailedEvent{
+		Failure:              errmsg,
+		CommandFinishedEvent: finishedEvent,
+	}
+
+	c.cmdMonitor.Failed(failureEvent)
+	return nil
+}
+
 func (c *connection) WriteWireMessage(ctx context.Context, wm wiremessage.WireMessage) error {
 	var err error
 	if c.dead {
@@ -415,6 +639,10 @@ func (c *connection) WriteWireMessage(ctx context.Context, wm wiremessage.WireMe
 	}
 
 	c.bumpIdleDeadline()
+	err = c.commandStartedEvent(wm)
+	if err != nil {
+		return err
+	}
 	return nil
 }
 
@@ -562,6 +790,11 @@ func (c *connection) ReadWireMessage(ctx context.Context) (wiremessage.WireMessa
 	}
 
 	c.bumpIdleDeadline()
+	err = c.commandFinishedEvent(wm)
+	if err != nil {
+		return nil, err // TODO: do we care if monitoring fails?
+	}
+
 	return wm, nil
 }
 

core/connection/options.go

@@ -11,6 +11,7 @@ import (
 	"time"
 
 	"github.com/mongodb/mongo-go-driver/core/compressor"
+	"github.com/mongodb/mongo-go-driver/core/event"
 )
 
 type config struct {
@@ -20,6 +21,7 @@ type config struct {
 	handshaker     Handshaker
 	idleTimeout    time.Duration
 	lifeTimeout    time.Duration
+	cmdMonitor     *event.CommandMonitor
 	readTimeout    time.Duration
 	writeTimeout   time.Duration
 	tlsConfig      *TLSConfig
@@ -133,3 +135,11 @@ func WithTLSConfig(fn func(*TLSConfig) *TLSConfig) Option {
 		return nil
 	}
 }
+
+// WithMonitor configures a event for command monitoring.
+func WithMonitor(fn func(*event.CommandMonitor) *event.CommandMonitor) Option {
+	return func(c *config) error {
+		c.cmdMonitor = fn(c.cmdMonitor)
+		return nil
+	}
+}

core/event/monitoring.go

@@ -0,0 +1,64 @@
+package event
+
+import (
+	"time"
+
+	"github.com/mongodb/mongo-go-driver/bson"
+)
+
+// CommandMetadata contains metadata about a command sent to the server.
+type CommandMetadata struct {
+	Name string
+	Time time.Time
+}
+
+// CreateMetadata creates metadata for a command.
+func CreateMetadata(name string) *CommandMetadata {
+	return &CommandMetadata{
+		Name: name,
+		Time: time.Now(),
+	}
+}
+
+// TimeDifference returns the difference between now and the time a command was sent in nanoseconds.
+func (cm *CommandMetadata) TimeDifference() int64 {
+	t := time.Now()
+	duration := t.Sub(cm.Time)
+	return duration.Nanoseconds()
+}
+
+// CommandStartedEvent represents an event generated when a command is sent to a server.
+type CommandStartedEvent struct {
+	Command      *bson.Document
+	DatabaseName string
+	CommandName  string
+	RequestID    int64
+	ConnectionID string
+}
+
+// CommandFinishedEvent represents a generic command finishing.
+type CommandFinishedEvent struct {
+	DurationNanos int64
+	CommandName   string
+	RequestID     int64
+	ConnectionID  string
+}
+
+// CommandSucceededEvent represents an event generated when a command's execution succeeds.
+type CommandSucceededEvent struct {
+	CommandFinishedEvent
+	Reply *bson.Document
+}
+
+// CommandFailedEvent represents an event generated when a command's execution fails.
+type CommandFailedEvent struct {
+	CommandFinishedEvent
+	Failure string
+}
+
+// CommandMonitor represents a monitor that is triggered for different events.
+type CommandMonitor struct {
+	Started   func(*CommandStartedEvent)
+	Succeeded func(*CommandSucceededEvent)
+	Failed    func(*CommandFailedEvent)
+}

core/topology/cursor.go

@@ -44,6 +44,7 @@ func newCursor(result bson.Reader, server *Server, opts ...option.CursorOptioner
 	c := &cursor{
 		current: -1,
 		server:  server,
+		opts:    opts,
 	}
 	var ok bool
 	for itr.Next() {