Go 语言 mgo 库中并发批量 Upsert MongoDB 文档的优化实践

package main

import (
    "context"
    "fmt"
    "log"
    "sync"
    "time"

    "gopkg.in/mgo.v2"
    "gopkg.in/mgo.v2/bson"
)

// MyDocument 定义了MongoDB文档的结构
type MyDocument struct {
    ID    bson.ObjectId `bson:"_id,omitempty"`
    Key   string        `bson:"key"`
    Value string        `bson:"value"`
    Count int           `bson:"count"`
}

// upsertDocument 函数用于执行单个文档的Upsert操作
// 注意：传入的session是主session，函数内部会进行Copy
func upsertDocument(s *mgo.Session, collection *mgo.Collection, doc MyDocument) error {
    // 在并发场景下，每个 goroutine 应该使用 session 的一个副本
    // 这样可以安全地共享连接池，而不会影响其他 goroutine 的操作
    session := s.Copy() // 复制会话，以便并发安全地使用连接池
    defer session.Close() // 确保会话在使用完毕后关闭

    // 定义查询条件：根据Key字段查找文档
    selector := bson.M{"key": doc.Key}
    // 定义更新内容：设置Value和Count字段，如果插入新文档则设置Key
    change := bson.M{
        "$set": bson.M{
            "value": doc.Value,
            "count": doc.Count,
        },
        "$setOnInsert": bson.M{ // 如果是插入操作，设置Key字段
            "key": doc.Key,
        },
    }

    // 执行 Upsert 操作
    _, err := collection.With(session).Upsert(selector, change)
    if err != nil {
        return fmt.Errorf("upsert document with key %s failed: %w", doc.Key, err)
    }
    return nil
}

func main() {
    // MongoDB 连接字符串
    mongoURI := "mongodb://localhost:27017" // 根据实际情况修改

    // 连接到 MongoDB
    session, err := mgo.Dial(mongoURI)
    if err != nil {
        log.Fatalf("Failed to connect to MongoDB: %v", err)
    }
    defer session.Close() // 确保主会话在程序结束时关闭

    // 设置会话模式，例如 Monotonic 或 Strong
    // Monotonic 模式在大多数读写分离场景下提供一致性保证，性能较好
    session.SetMode(mgo.Monotonic, true)

    // 获取集合
    collection := session.DB("testdb").C("mydocuments")

    // 清空集合以便测试 (可选)
    // if err := collection.DropCollection(); err != nil {
    //  log.Printf("Failed to drop collection: %v", err)
    // }

    // 准备要 Upsert 的文档数据
    documentsToUpsert := []MyDocument{
        {Key: "doc1", Value: "initial value 1", Count: 1},
        {Key: "doc2", Value: "initial value 2", Count: 2},
        {Key: "doc3", Value: "initial value 3", Count: 3},
        {Key: "doc1", Value: "updated value 1", Count: 10}, // 更新 doc1
        {Key: "doc4", Value: "new value 4", Count: 4},
        {Key: "doc5", Value: "new value 5", Count: 5},
        {Key: "doc2", Value: "updated value 2", Count: 20}, // 更新 doc2
        {Key: "doc6", Value: "new value 6", Count: 6},
    }

    var wg sync.WaitGroup // 用于等待所有 goroutine 完成
    errCh := make(chan error, len(documentsToUpsert)) // 有缓冲通道，用于收集并发错误

    start := time.Now() // 记录开始时间

    fmt.Printf("开始并发 Upsert %d 个文档...\n", len(documentsToUpsert))

    // 遍历文档数据，为每个文档启动一个 goroutine
    for _, doc := range documentsToUpsert {
        wg.Add(1) // 增加 WaitGroup 计数
        go func(d MyDocument) {
            defer wg.Done() // goroutine 完成时减少 WaitGroup 计数
            // 调用 upsertDocument 函数执行 Upsert
            if err := upsertDocument(session, collection, d); err != nil {
                errCh <- err // 如果发生错误，发送到错误通道
            }
        }(doc) // 将当前文档作为参数传递给 goroutine
    }

    wg.Wait() // 等待所有 goroutine 完成
    close(errCh) // 关闭错误通道，以便安全地遍历

    // 检查是否有错误发生
    hasErrors := false
    for err := range errCh {
        log.Printf("并发 Upsert 错误: %v", err)
        hasErrors = true
    }

    if hasErrors {
        fmt.Println("部分或全部文档 Upsert 失败。")
    } else {
        fmt.Println("所有文档成功并发 Upsert。")
    }

    duration := time.Since(start) // 计算总耗时
    fmt.Printf("并发 Upsert 完成，耗时: %s\n", duration)

    // 验证数据库中的文档
    fmt.Println("\n验证数据库中的文档:")
    var results []MyDocument
    err = collection.Find(nil).All(&results)
    if err != nil {
        log.Fatalf("Failed to find documents: %v", err)
    }
    for _, res := range results {
        fmt.Printf("  Key: %s, Value: %s, Count: %d\n", res.Key, res.Value, res.Count)
    }
}