Replace bg reads with fg reads

Author: prestonvasquez

event/monitoring.go

@@ -91,24 +91,29 @@ const (
 
 // strings for pool command monitoring types
 const (
-	PoolCreated        = "ConnectionPoolCreated"
-	PoolReady          = "ConnectionPoolReady"
-	PoolCleared        = "ConnectionPoolCleared"
-	PoolClosedEvent    = "ConnectionPoolClosed"
-	ConnectionCreated  = "ConnectionCreated"
-	ConnectionReady    = "ConnectionReady"
-	ConnectionClosed   = "ConnectionClosed"
-	GetStarted         = "ConnectionCheckOutStarted"
-	GetFailed          = "ConnectionCheckOutFailed"
-	GetSucceeded       = "ConnectionCheckedOut"
-	ConnectionReturned = "ConnectionCheckedIn"
+	PoolCreated                    = "ConnectionPoolCreated"
+	PoolReady                      = "ConnectionPoolReady"
+	PoolCleared                    = "ConnectionPoolCleared"
+	PoolClosedEvent                = "ConnectionPoolClosed"
+	ConnectionCreated              = "ConnectionCreated"
+	ConnectionReady                = "ConnectionReady"
+	ConnectionClosed               = "ConnectionClosed"
+	GetStarted                     = "ConnectionCheckOutStarted"
+	GetFailed                      = "ConnectionCheckOutFailed"
+	GetSucceeded                   = "ConnectionCheckedOut"
+	ConnectionReturned             = "ConnectionCheckedIn"
+	ConnectionPendingReadStarted   = "ConnectionPendingReadStarted"
+	ConnectionPendingReadSucceeded = "ConnectionPendingReadSucceeded"
+	ConnectionPendingReadFailed    = "ConnectionPendingReadFailed"
 )
 
 // MonitorPoolOptions contains pool options as formatted in pool events
 type MonitorPoolOptions struct {
-	MaxPoolSize        uint64 `json:"maxPoolSize"`
-	MinPoolSize        uint64 `json:"minPoolSize"`
-	WaitQueueTimeoutMS uint64 `json:"maxIdleTimeMS"`
+	MaxPoolSize        uint64        `json:"maxPoolSize"`
+	MinPoolSize        uint64        `json:"minPoolSize"`
+	WaitQueueTimeoutMS uint64        `json:"maxIdleTimeMS"`
+	RequestID          int32         `json:"requestId"`
+	RemainingTime      time.Duration `json:"remainingTime"`
 }
 
 // PoolEvent contains all information summarizing a pool event

internal/driverutil/context.go

@@ -0,0 +1,28 @@
+package driverutil
+
+import "context"
+
+type ContextKey string
+
+const (
+	ContextKeyHasMaxTimeMS ContextKey = "hasMaxTimeMS"
+	ContextKeyRequestID    ContextKey = "requestID"
+)
+
+func WithValueHasMaxTimeMS(parentCtx context.Context, val bool) context.Context {
+	return context.WithValue(parentCtx, ContextKeyHasMaxTimeMS, val)
+}
+
+func WithRequestID(parentCtx context.Context, requestID int32) context.Context {
+	return context.WithValue(parentCtx, ContextKeyRequestID, requestID)
+}
+
+func HasMaxTimeMS(ctx context.Context) bool {
+	return ctx.Value(ContextKeyHasMaxTimeMS) != nil
+}
+
+func GetRequestID(ctx context.Context) (int32, bool) {
+	val, ok := ctx.Value(ContextKeyRequestID).(int32)
+
+	return val, ok
+}

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"

internal/ptrutil/ptr.go

@@ -0,0 +1,12 @@
+// Copyright (C) MongoDB, Inc. 2024-present.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"); you may
+// not use this file except in compliance with the License. You may obtain
+// a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
+
+package ptrutil
+
+// Ptr will return the memory location of the given value.
+func Ptr[T any](val T) *T {
+	return &val
+}

x/mongo/driver/operation.go

@@ -807,6 +807,12 @@ func (op Operation) Execute(ctx context.Context) error {
 			if moreToCome {
 				roundTrip = op.moreToComeRoundTrip
 			}
+
+			if maxTimeMS != 0 {
+				ctx = driverutil.WithValueHasMaxTimeMS(ctx, true)
+				ctx = driverutil.WithRequestID(ctx, startedInfo.requestID)
+			}
+
 			res, err = roundTrip(ctx, conn, *wm)
 
 			if ep, ok := srvr.(ErrorProcessor); ok {

x/mongo/driver/topology/connection.go

@@ -21,6 +21,7 @@ import (
 	"time"
 
 	"go.mongodb.org/mongo-driver/internal/csot"
+	"go.mongodb.org/mongo-driver/internal/driverutil"
 	"go.mongodb.org/mongo-driver/mongo/address"
 	"go.mongodb.org/mongo-driver/mongo/description"
 	"go.mongodb.org/mongo-driver/x/bsonx/bsoncore"
@@ -46,6 +47,12 @@ var (
 
 func nextConnectionID() uint64 { return atomic.AddUint64(&globalConnectionID, 1) }
 
+type pendingReadState struct {
+	remainingBytes int32
+	requestID      int32
+	remainingTime  *time.Duration
+}
+
 type connection struct {
 	// state must be accessed using the atomic package and should be at the beginning of the struct.
 	// - atomic bug: https://pkg.go.dev/sync/atomic#pkg-note-BUG
@@ -83,7 +90,8 @@ type connection struct {
 
 	// awaitRemainingBytes indicates the size of server response that was not completely
 	// read before returning the connection to the pool.
-	awaitRemainingBytes *int32
+	// awaitRemainingBytes *int32
+	pendingReadState *pendingReadState
 
 	// oidcTokenGenID is the monotonic generation ID for OIDC tokens, used to invalidate
 	// accessTokens in the OIDC authenticator cache.
@@ -452,7 +460,7 @@ func (c *connection) readWireMessage(ctx context.Context) ([]byte, error) {
 
 	dst, errMsg, err := c.read(ctx)
 	if err != nil {
-		if c.awaitRemainingBytes == nil {
+		if c.pendingReadState == nil {
 			// If the connection was not marked as awaiting response, use the
 			// pre-CSOT behavior and close the connection because we don't know
 			// if there are other bytes left to read.
@@ -523,8 +531,13 @@ 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) {
-			c.awaitRemainingBytes = &l
+		if l := int32(n); l == 0 && isCSOTTimeout(err) && driverutil.HasMaxTimeMS(ctx) {
+			requestID, _ := driverutil.GetRequestID(ctx)
+
+			c.pendingReadState = &pendingReadState{
+				remainingBytes: l,
+				requestID:      requestID,
+			}
 		}
 		return nil, "incomplete read of message header", err
 	}
@@ -539,8 +552,13 @@ 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) {
-			c.awaitRemainingBytes = &remainingBytes
+		if remainingBytes > 0 && isCSOTTimeout(err) && driverutil.HasMaxTimeMS(ctx) {
+			requestID, _ := driverutil.GetRequestID(ctx)
+
+			c.pendingReadState = &pendingReadState{
+				remainingBytes: remainingBytes,
+				requestID:      requestID,
+			}
 		}
 		return dst, "incomplete read of full message", err
 	}