Advance logic to meet TD

Author: prestonvasquez

event/monitoring.go

@@ -108,9 +108,11 @@ type PoolEvent struct {
 	Reason       string              `json:"reason"`
 	// ServiceID is only set if the Type is PoolCleared and the server is deployed behind a load balancer. This field
 	// can be used to distinguish between individual servers in a load balanced deployment.
-	ServiceID    *bson.ObjectID `json:"serviceId"`
-	Interruption bool           `json:"interruptInUseConnections"`
-	Error        error          `json:"error"`
+	ServiceID     *bson.ObjectID `json:"serviceId"`
+	Interruption  bool           `json:"interruptInUseConnections"`
+	Error         error          `json:"error"`
+	RequestID     int32          `json:"requestId"`
+	RemainingTime time.Duration  `json:"remainingTime"`
 }
 
 // PoolMonitor is a function that allows the user to gain access to events occurring in the pool

internal/logger/component.go

@@ -28,6 +28,9 @@ const (
 	ConnectionCheckoutFailed         = "Connection checkout failed"
 	ConnectionCheckedOut             = "Connection checked out"
 	ConnectionCheckedIn              = "Connection checked in"
+	ConnectionPendingReadStarted     = "Pending read started"
+	ConnectionPendingReadSucceeded   = "Pending read succeeded"
+	ConnectionPendingReadFailed      = "Pending read failed"
 	ServerSelectionFailed            = "Server selection failed"
 	ServerSelectionStarted           = "Server selection started"
 	ServerSelectionSucceeded         = "Server selection succeeded"

x/mongo/driver/drivertest/channel_conn.go

@@ -13,6 +13,7 @@ import (
 	"go.mongodb.org/mongo-driver/v2/mongo/address"
 	"go.mongodb.org/mongo-driver/v2/x/bsonx/bsoncore"
 	"go.mongodb.org/mongo-driver/v2/x/mongo/driver/description"
+	"go.mongodb.org/mongo-driver/v2/x/mongo/driver/mnet"
 	"go.mongodb.org/mongo-driver/v2/x/mongo/driver/wiremessage"
 )
 
@@ -52,7 +53,7 @@ func (c *ChannelConn) Write(ctx context.Context, wm []byte) error {
 }
 
 // ReadWireMessage implements the driver.Connection interface.
-func (c *ChannelConn) Read(ctx context.Context) ([]byte, error) {
+func (c *ChannelConn) Read(ctx context.Context, _ ...mnet.ReadOption) ([]byte, error) {
 	var wm []byte
 	var err error
 	select {

x/mongo/driver/drivertest/opmsg_deployment.go

@@ -68,7 +68,7 @@ func (c *connection) SetOIDCTokenGenID(uint64) {
 }
 
 // Read returns the next response in the connection's list of responses.
-func (c *connection) Read(_ context.Context) ([]byte, error) {
+func (c *connection) Read(_ context.Context, _ ...mnet.ReadOption) ([]byte, error) {
 	var dst []byte
 	if len(c.responses) == 0 {
 		return dst, errors.New("no responses remaining")

x/mongo/driver/mnet/connection.go

@@ -14,10 +14,29 @@ import (
 	"go.mongodb.org/mongo-driver/v2/x/mongo/driver/description"
 )
 
+type ReadOption func(*ReadOptions)
+
+type ReadOptions struct {
+	HasMaxTimeMS bool
+	RequestID    int32
+}
+
+func WithReadMaxTimeMS() ReadOption {
+	return func(opts *ReadOptions) {
+		opts.HasMaxTimeMS = true
+	}
+}
+
+func WithRequestID(requestID int32) ReadOption {
+	return func(opts *ReadOptions) {
+		opts.RequestID = requestID
+	}
+}
+
 // ReadWriteCloser represents a Connection where server operations
 // can read from, written to, and closed.
 type ReadWriteCloser interface {
-	Read(ctx context.Context) ([]byte, error)
+	Read(ctx context.Context, opts ...ReadOption) ([]byte, error)
 	Write(ctx context.Context, wm []byte) error
 	io.Closer
 }

x/mongo/driver/operation.go

@@ -792,7 +792,18 @@ func (op Operation) Execute(ctx context.Context) error {
 			if moreToCome {
 				roundTrip = op.moreToComeRoundTrip
 			}
-			res, err = roundTrip(ctx, conn, *wm)
+
+			readOpts := []mnet.ReadOption{}
+			if maxTimeMS != 0 {
+				readOpts = append(readOpts, mnet.WithReadMaxTimeMS())
+				readOpts = append(readOpts, mnet.WithRequestID(startedInfo.requestID))
+			}
+
+			// Inform the roundTrip if maxTimeMS is set. If it is and the operation
+			// times out, then the connection should be put into a "pending" state
+			// so that the next time it is checked out it attempts to finish the read
+			// which is almost certainly a server error noting a timeout.
+			res, err = roundTrip(ctx, conn, *wm, readOpts)
 
 			if ep, ok := srvr.(ErrorProcessor); ok {
 				_ = ep.ProcessError(err, conn)
@@ -1076,16 +1087,25 @@ func (op Operation) retryable(desc description.Server) bool {
 
 // roundTrip writes a wiremessage to the connection and then reads a wiremessage. The wm parameter
 // is reused when reading the wiremessage.
-func (op Operation) roundTrip(ctx context.Context, conn *mnet.Connection, wm []byte) ([]byte, error) {
+func (op Operation) roundTrip(
+	ctx context.Context,
+	conn *mnet.Connection,
+	wm []byte,
+	readOpts []mnet.ReadOption,
+) ([]byte, error) {
 	err := conn.Write(ctx, wm)
 	if err != nil {
 		return nil, op.networkError(err)
 	}
-	return op.readWireMessage(ctx, conn)
+	return op.readWireMessage(ctx, conn, readOpts...)
 }
 
-func (op Operation) readWireMessage(ctx context.Context, conn *mnet.Connection) (result []byte, err error) {
-	wm, err := conn.Read(ctx)
+func (op Operation) readWireMessage(
+	ctx context.Context,
+	conn *mnet.Connection,
+	opts ...mnet.ReadOption,
+) (result []byte, err error) {
+	wm, err := conn.Read(ctx, opts...)
 	if err != nil {
 		return nil, op.networkError(err)
 	}
@@ -1156,7 +1176,12 @@ func (op Operation) networkError(err error) error {
 
 // moreToComeRoundTrip writes a wiremessage to the provided connection. This is used when an OP_MSG is
 // being sent with  the moreToCome bit set.
-func (op *Operation) moreToComeRoundTrip(ctx context.Context, conn *mnet.Connection, wm []byte) (result []byte, err error) {
+func (op *Operation) moreToComeRoundTrip(
+	ctx context.Context,
+	conn *mnet.Connection,
+	wm []byte,
+	_ []mnet.ReadOption,
+) (result []byte, err error) {
 	err = conn.Write(ctx, wm)
 	if err != nil {
 		if op.Client != nil {

x/mongo/driver/operation_test.go

@@ -793,7 +793,7 @@ func (m *mockConnection) Write(_ context.Context, wm []byte) error {
 	return m.rWriteErr
 }
 
-func (m *mockConnection) Read(_ context.Context) ([]byte, error) {
+func (m *mockConnection) Read(_ context.Context, _ ...mnet.ReadOption) ([]byte, error) {
 	return m.rReadWM, m.rReadErr
 }
 

x/mongo/driver/topology/connection.go

@@ -86,6 +86,7 @@ type connection struct {
 	// awaitRemainingBytes indicates the size of server response that was not completely
 	// read before returning the connection to the pool.
 	awaitRemainingBytes *int32
+	requestID           int32
 	remainingTime       *time.Duration
 	pendingReadMU       sync.Mutex
 }
@@ -396,7 +397,7 @@ func (c *connection) write(ctx context.Context, wm []byte) (err error) {
 }
 
 // readWireMessage reads a wiremessage from the connection. The dst parameter will be overwritten.
-func (c *connection) readWireMessage(ctx context.Context) ([]byte, error) {
+func (c *connection) readWireMessage(ctx context.Context, opts ...mnet.ReadOption) ([]byte, error) {
 	if atomic.LoadInt64(&c.state) != connConnected {
 		return nil, ConnectionError{
 			ConnectionID: c.id,
@@ -409,7 +410,7 @@ func (c *connection) readWireMessage(ctx context.Context) ([]byte, error) {
 		return nil, ConnectionError{ConnectionID: c.id, Wrapped: err, message: "failed to set read deadline"}
 	}
 
-	dst, errMsg, err := c.read(ctx)
+	dst, errMsg, err := c.read(ctx, opts...)
 	if err != nil {
 		c.pendingReadMU.Lock()
 		if c.awaitRemainingBytes == nil {
@@ -452,7 +453,7 @@ func (c *connection) parseWmSizeBytes(wmSizeBytes [4]byte) (int32, error) {
 	return size, nil
 }
 
-func (c *connection) read(ctx context.Context) (bytesRead []byte, errMsg string, err error) {
+func (c *connection) read(ctx context.Context, opts ...mnet.ReadOption) (bytesRead []byte, errMsg string, err error) {
 	go c.cancellationListener.Listen(ctx, c.cancellationListenerCallback)
 	defer func() {
 		// If the context is cancelled after we finish reading the server response, the cancellation listener could fire
@@ -465,6 +466,11 @@ func (c *connection) read(ctx context.Context) (bytesRead []byte, errMsg string,
 		}
 	}()
 
+	readOpts := mnet.ReadOptions{}
+	for _, opt := range opts {
+		opt(&readOpts)
+	}
+
 	isCSOTTimeout := func(err error) bool {
 		// If the error was a timeout error, instead of closing the
 		// connection mark it as awaiting response so the pool can read the
@@ -482,9 +488,10 @@ func (c *connection) read(ctx context.Context) (bytesRead []byte, errMsg string,
 	// reading messages from an exhaust cursor.
 	n, err := io.ReadFull(c.nc, sizeBuf[:])
 	if err != nil {
-		if l := int32(n); l == 0 && isCSOTTimeout(err) {
+		if l := int32(n); l == 0 && isCSOTTimeout(err) && readOpts.HasMaxTimeMS {
 			c.pendingReadMU.Lock()
 			c.awaitRemainingBytes = &l
+			c.requestID = readOpts.RequestID
 			c.remainingTime = ptrutil.Ptr(PendingReadTimeout)
 			c.pendingReadMU.Unlock()
 		}
@@ -501,9 +508,10 @@ func (c *connection) read(ctx context.Context) (bytesRead []byte, errMsg string,
 	n, err = io.ReadFull(c.nc, dst[4:])
 	if err != nil {
 		remainingBytes := size - 4 - int32(n)
-		if remainingBytes > 0 && isCSOTTimeout(err) {
+		if remainingBytes > 0 && isCSOTTimeout(err) && readOpts.HasMaxTimeMS {
 			c.pendingReadMU.Lock()
 			c.awaitRemainingBytes = &remainingBytes
+			c.requestID = readOpts.RequestID
 			c.remainingTime = ptrutil.Ptr(PendingReadTimeout)
 			c.pendingReadMU.Unlock()
 		}
@@ -664,8 +672,8 @@ func (c initConnection) LocalAddress() address.Address {
 func (c initConnection) Write(ctx context.Context, wm []byte) error {
 	return c.writeWireMessage(ctx, wm)
 }
-func (c initConnection) Read(ctx context.Context) ([]byte, error) {
-	return c.readWireMessage(ctx)
+func (c initConnection) Read(ctx context.Context, opts ...mnet.ReadOption) ([]byte, error) {
+	return c.readWireMessage(ctx, opts...)
 }
 func (c initConnection) SetStreaming(streaming bool) {
 	c.setStreaming(streaming)
@@ -712,13 +720,13 @@ func (c *Connection) Write(ctx context.Context, wm []byte) error {
 
 // ReadWireMessage handles reading a wire message from the underlying connection. The dst parameter
 // will be overwritten with the new wire message.
-func (c *Connection) Read(ctx context.Context) ([]byte, error) {
+func (c *Connection) Read(ctx context.Context, opts ...mnet.ReadOption) ([]byte, error) {
 	c.mu.RLock()
 	defer c.mu.RUnlock()
 	if c.connection == nil {
 		return nil, ErrConnectionClosed
 	}
-	return c.connection.readWireMessage(ctx)
+	return c.connection.readWireMessage(ctx, opts...)
 }
 
 // CompressWireMessage handles compressing the provided wire message using the underlying

x/mongo/driver/topology/pool.go

@@ -791,7 +791,7 @@ func (p *pool) removeConnection(conn *connection, reason reason, err error) erro
 // Deprecated: PendingReadTimeout is intended for internal use only and may be
 // removed or modified at any time.
 
-var PendingReadTimeout = 1 * time.Second
+var PendingReadTimeout = 400 * time.Millisecond
 
 // awaitPendingRead sets a new read deadline on the provided connection and
 // tries to read any bytes returned by the server. If there are any errors, the
@@ -805,18 +805,32 @@ func awaitPendingRead(ctx context.Context, pool *pool, conn *connection) error {
 		return nil
 	}
 
-	//if pool.monitor != nil {
-	//	pool.monitor.Event(&event.PoolEvent{
-	//		Type:         event.ConnectionPendingReadStarted,
-	//		ConnectionID: conn.driverConnectionID,
-	//	})
-	//}
+	if mustLogPoolMessage(pool) {
+		keysAndValues := logger.KeyValues{
+			logger.KeyDriverConnectionID, conn.driverConnectionID,
+			logger.KeyRequestID, conn.requestID,
+		}
+
+		logPoolMessage(pool, logger.ConnectionPendingReadStarted, keysAndValues...)
+	}
 
 	size := *conn.awaitRemainingBytes
 
 	checkIn := false
+	var someErr error
 
 	defer func() {
+		if mustLogPoolMessage(pool) && someErr != nil {
+			keysAndValues := logger.KeyValues{
+				logger.KeyDriverConnectionID, conn.driverConnectionID,
+				logger.KeyRequestID, conn.requestID,
+				logger.KeyReason, someErr.Error(),
+				logger.KeyRemainingTimeMS, *conn.remainingTime,
+			}
+
+			logPoolMessage(pool, logger.ConnectionPendingReadFailed, keysAndValues...)
+		}
+
 		// If we have exceeded the time limit, then close the connection.
 		if conn.remainingTime != nil && *conn.remainingTime < 0 {
 			if err := conn.close(); err != nil {
@@ -830,14 +844,6 @@ func awaitPendingRead(ctx context.Context, pool *pool, conn *connection) error {
 			return
 		}
 
-		//if pool.monitor != nil {
-		//	pool.monitor.Event(&event.PoolEvent{
-		//		Type:         event.ConnectionPendingReadFailed,
-		//		ConnectionID: conn.driverConnectionID,
-		//		//Reason:       readErr.Error(),
-		//	})
-		//}
-
 		// No matter what happens, always check the connection back into the
 		// pool, which will either make it available for other operations or
 		// remove it from the pool if it was closed.
@@ -857,7 +863,10 @@ func awaitPendingRead(ctx context.Context, pool *pool, conn *connection) error {
 	err := conn.nc.SetReadDeadline(dl)
 	if err != nil {
 		checkIn = true
-		return fmt.Errorf("error setting a read deadline: %w", err)
+
+		someErr = fmt.Errorf("error setting a read deadline: %w", err)
+
+		return someErr
 	}
 
 	st := time.Now()
@@ -869,15 +878,16 @@ func awaitPendingRead(ctx context.Context, pool *pool, conn *connection) error {
 			checkIn = true
 
 			err = transformNetworkError(ctx, err, contextDeadlineUsed)
+			someErr = fmt.Errorf("error reading the message size: %w", err)
 
-			return fmt.Errorf("error reading the message size: %w", err)
+			return someErr
 		}
 		size, err = conn.parseWmSizeBytes(sizeBuf)
 		if err != nil {
 			checkIn = true
-			err = transformNetworkError(ctx, err, contextDeadlineUsed)
+			someErr = transformNetworkError(ctx, err, contextDeadlineUsed)
 
-			return err
+			return someErr
 		}
 		size -= 4
 	}
@@ -893,17 +903,19 @@ func awaitPendingRead(ctx context.Context, pool *pool, conn *connection) error {
 		checkIn = true
 
 		err = transformNetworkError(ctx, err, contextDeadlineUsed)
+		someErr = fmt.Errorf("error discarding %d byte message: %w", size, err)
 
-		return fmt.Errorf("error discarding %d byte message: %w", size, err)
+		return someErr
 	}
 
-	//if pool.monitor != nil {
-	//	pool.monitor.Event(&event.PoolEvent{
-	//		Type:         event.ConnectionPendingReadSucceeded,
-	//		ConnectionID: conn.driverConnectionID,
-	//		//Reason:       readErr.Error(),
-	//	})
-	//}
+	if mustLogPoolMessage(pool) {
+		keysAndValues := logger.KeyValues{
+			logger.KeyDriverConnectionID, conn.driverConnectionID,
+			logger.KeyRequestID, conn.requestID,
+		}
+
+		logPoolMessage(pool, logger.ConnectionPendingReadSucceeded, keysAndValues...)
+	}
 
 	conn.awaitRemainingBytes = nil
 	conn.remainingTime = nil