Skip to content

feat: Add MongoDB read preference support #92

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
scmsscm wants to merge 1 commit into
base: master
Choose a base branch
from
Open

feat: Add MongoDB read preference support #92

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 9 additions & 4 deletions mongo/mongo.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -154,7 +154,8 @@ func (m *Mongo) RoundTrip(msg *Message, tags []string) (_ *Message, err error) {
requestCursorID, _ := msg.Op.CursorID()
requestCommand, collection := msg.Op.CommandAndCollection()
transactionDetails := msg.Op.TransactionDetails()
server, err := m.selectServer(requestCursorID, collection, transactionDetails)
readPref, _ := msg.Op.ReadPref()
server, err := m.selectServer(requestCursorID, collection, transactionDetails, readPref)
if err != nil {
return nil, err
}
Expand Down Expand Up @@ -230,7 +231,7 @@ func (m *Mongo) RoundTrip(msg *Message, tags []string) (_ *Message, err error) {
}, nil
}

func (m *Mongo) selectServer(requestCursorID int64, collection string, transDetails *TransactionDetails) (server driver.Server, err error) {
func (m *Mongo) selectServer(requestCursorID int64, collection string, transDetails *TransactionDetails, rp *readpref.ReadPref) (server driver.Server, err error) {
defer func(start time.Time) {
_ = m.statsd.Timing("server_selection", time.Since(start), []string{fmt.Sprintf("success:%v", err == nil)}, 1)
}(time.Now())
Expand All @@ -252,9 +253,13 @@ func (m *Mongo) selectServer(requestCursorID int64, collection string, transDeta
}
}

// Select a server
// Select a server using the provided read preference
// If no read preference was provided, default to primary
if rp == nil {
rp = readpref.Primary()
}
selector := description.CompositeSelector([]description.ServerSelector{
description.ReadPrefSelector(readpref.Primary()), // ignored by sharded clusters
description.ReadPrefSelector(rp), // use client's read preference (ignored by sharded clusters)
description.LatencySelector(15 * time.Millisecond), // default localThreshold for the client
})
return m.topology.SelectServer(m.roundTripCtx, selector)
Expand Down
97 changes: 97 additions & 0 deletions mongo/operations.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ import (
"fmt"
"strings"

"go.mongodb.org/mongo-driver/mongo/readpref"
"go.mongodb.org/mongo-driver/x/bsonx/bsoncore"
"go.mongodb.org/mongo-driver/x/mongo/driver"
"go.mongodb.org/mongo-driver/x/mongo/driver/wiremessage"
Expand Down Expand Up @@ -33,6 +34,7 @@ type Operation interface {
Unacknowledged() bool
CommandAndCollection() (Command, string)
TransactionDetails() *TransactionDetails
ReadPref() (rp *readpref.ReadPref, ok bool)
}

// see https://github.com/mongodb/mongo-go-driver/blob/v1.7.2/x/mongo/driver/operation.go#L1361-L1426
Expand Down Expand Up @@ -121,6 +123,10 @@ func (o *opUnknown) String() string {
return fmt.Sprintf("{ OpUnknown opCode: %d, wm: %s }", o.opCode, o.wm)
}

func (o *opUnknown) ReadPref() (rp *readpref.ReadPref, ok bool) {
return readpref.Primary(), false
}

// https://docs.mongodb.com/manual/reference/mongodb-wire-protocol/#wire-op-query
type opQuery struct {
reqID int32
Expand Down Expand Up @@ -230,6 +236,10 @@ func (q *opQuery) String() string {
return fmt.Sprintf("{ OpQuery flags: %s, fullCollectionName: %s, numberToSkip: %d, numberToReturn: %d, query: %s, returnFieldsSelector: %s }", q.flags.String(), q.fullCollectionName, q.numberToSkip, q.numberToReturn, q.query.String(), q.returnFieldsSelector.String())
}

func (q *opQuery) ReadPref() (rp *readpref.ReadPref, ok bool) {
return extractReadPref(q.query)
}

// https://docs.mongodb.com/manual/reference/mongodb-wire-protocol/#op-msg
type opMsg struct {
reqID int32
Expand Down Expand Up @@ -485,6 +495,17 @@ func (m *opMsg) String() string {
return fmt.Sprintf("{ OpMsg flags: %d, sections: [%s], checksum: %d }", m.flags, strings.Join(sections, ", "), m.checksum)
}

func (m *opMsg) ReadPref() (rp *readpref.ReadPref, ok bool) {
if len(m.sections) == 0 {
return readpref.Primary(), false
}
single, ok := m.sections[0].(*opMsgSectionSingle)
if !ok {
return readpref.Primary(), false
}
return extractReadPref(single.msg)
}

// https://docs.mongodb.com/manual/reference/mongodb-wire-protocol/#op-reply
type opReply struct {
reqID int32
Expand Down Expand Up @@ -588,6 +609,11 @@ func (r *opReply) String() string {
return fmt.Sprintf("{ OpReply flags: %d, cursorID: %d, startingFrom: %d, numReturned: %d, documents: [%s] }", r.flags, r.cursorID, r.startingFrom, r.numReturned, strings.Join(documents, ", "))
}

func (r *opReply) ReadPref() (rp *readpref.ReadPref, ok bool) {
// Replies don't contain read preferences
return readpref.Primary(), false
}

// https://docs.mongodb.com/manual/reference/mongodb-wire-protocol/#op-get-more
type opGetMore struct {
reqID int32
Expand Down Expand Up @@ -676,6 +702,11 @@ func (g *opGetMore) String() string {
return fmt.Sprintf("{ OpGetMore fullCollectionName: %s, numberToReturn: %d, cursorID: %d }", g.fullCollectionName, g.numberToReturn, g.cursorID)
}

func (g *opGetMore) ReadPref() (rp *readpref.ReadPref, ok bool) {
// GetMore operations don't specify read preference, they inherit from cursor
return readpref.Primary(), false
}

// https://docs.mongodb.com/manual/reference/mongodb-wire-protocol/#op_update
type opUpdate struct {
reqID int32
Expand Down Expand Up @@ -761,6 +792,11 @@ func (u *opUpdate) String() string {
return fmt.Sprintf("{ OpQuery fullCollectionName: %s, flags: %d, selector: %s, update: %s }", u.fullCollectionName, u.flags, u.selector.String(), u.update.String())
}

func (u *opUpdate) ReadPref() (rp *readpref.ReadPref, ok bool) {
// Update is a write operation, always goes to primary
return readpref.Primary(), false
}

// https://docs.mongodb.com/manual/reference/mongodb-wire-protocol/#op_insert
type opInsert struct {
reqID int32
Expand Down Expand Up @@ -845,6 +881,11 @@ func (i *opInsert) String() string {
return fmt.Sprintf("{ OpInsert flags: %d, fullCollectionName: %s, documents: %s }", i.flags, i.fullCollectionName, strings.Join(documents, ", "))
}

func (i *opInsert) ReadPref() (rp *readpref.ReadPref, ok bool) {
// Insert is a write operation, always goes to primary
return readpref.Primary(), false
}

// https://docs.mongodb.com/manual/reference/mongodb-wire-protocol/#op_insert
type opDelete struct {
reqID int32
Expand Down Expand Up @@ -928,6 +969,11 @@ func (d *opDelete) String() string {
return fmt.Sprintf("{ OpDelete fullCollectionName: %s, flags: %d, selector: %s }", d.fullCollectionName, d.flags, d.selector.String())
}

func (d *opDelete) ReadPref() (rp *readpref.ReadPref, ok bool) {
// Delete is a write operation, always goes to primary
return readpref.Primary(), false
}

// https://docs.mongodb.com/manual/reference/mongodb-wire-protocol/#op_kill_cursors
type opKillCursors struct {
reqID int32
Expand Down Expand Up @@ -1006,6 +1052,11 @@ func (k *opKillCursors) String() string {
return fmt.Sprintf("{ OpKillCursors cursorIDs: %v }", k.cursorIDs)
}

func (k *opKillCursors) ReadPref() (rp *readpref.ReadPref, ok bool) {
// KillCursors doesn't use read preferences
return readpref.Primary(), false
}

func appendi32(dst []byte, i32 int32) []byte {
return append(dst, byte(i32), byte(i32>>8), byte(i32>>16), byte(i32>>24))
}
Expand All @@ -1030,3 +1081,49 @@ func readCString(src []byte) (string, []byte, bool) {
}
return string(src[:idx]), src[idx+1:], true
}

// extractReadPref extracts the read preference from a BSON document
// MongoDB clients send read preferences in the $readPreference field
// See: https://github.com/mongodb/specifications/blob/master/source/server-selection/server-selection.rst#passing-read-preference-to-mongos
func extractReadPref(doc bsoncore.Document) (*readpref.ReadPref, bool) {
// Check for $readPreference field
rpVal := doc.Lookup("$readPreference")
if rpVal.Type == 0 {
// No read preference specified, default to primary
return readpref.Primary(), false
}

rpDoc, ok := rpVal.DocumentOK()
if !ok {
return readpref.Primary(), false
}

// Extract mode
modeVal := rpDoc.Lookup("mode")
mode, ok := modeVal.StringValueOK()
if !ok {
return readpref.Primary(), false
}

// Parse the mode string and create appropriate read preference
var rp *readpref.ReadPref
switch mode {
case "primary":
rp = readpref.Primary()
case "primaryPreferred":
rp = readpref.PrimaryPreferred()
case "secondary":
rp = readpref.Secondary()
case "secondaryPreferred":
rp = readpref.SecondaryPreferred()
case "nearest":
rp = readpref.Nearest()
default:
return readpref.Primary(), false
}

// TODO: Support tag sets and maxStalenessSeconds if needed
// For now, we just use the mode without additional options

return rp, true
}
Loading