Golang mgo库：多文档Upsert操作的并发优化策略与实践

package main

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

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

// MyDocument 定义文档结构
type MyDocument struct {
    ID        bson.ObjectId `bson:"_id,omitempty"` // MongoDB的_id字段
    Key       string        `bson:"key"`           // 用于Upsert的唯一键
    Value     string        `bson:"value"`         // 更新的值
    CreatedAt time.Time     `bson:"createdAt,omitempty"` // 创建时间
    UpdatedAt time.Time     `bson:"updatedAt"`           // 更新时间
}

func main() {
    // 1. 连接MongoDB
    session, err := mgo.Dial("mongodb://localhost:27017")
    if err != nil {
        log.Fatalf("连接MongoDB失败: %v", err)
    }
    defer session.Close() // 确保主会话在程序结束时关闭

    // 设置会话模式，例如 Monotonic (读写操作发生在同一连接上)
    session.SetMode(mgo.Monotonic, true)

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

    // 清理旧数据 (可选)
    // if _, err := collection.RemoveAll(nil); err != nil {
    //  log.Printf("清理旧数据失败: %v", err)
    // }

    // 2. 准备需要Upsert的数据
    docsToUpsert := []MyDocument{
        {Key: "productA", Value: "初始版本 A"},
        {Key: "productB", Value: "初始版本 B"},
        {Key: "productC", Value: "初始版本 C"},
        {Key: "productA", Value: "更新版本 A"}, // 再次Upsert productA
        {Key: "productD", Value: "新增产品 D"},
        {Key: "productB", Value: "二次更新 B"}, // 再次Upsert productB
    }

    var wg sync.WaitGroup
    // 使用带缓冲的通道来收集错误，防止goroutine阻塞
    errors := make(chan error, len(docsToUpsert))

    fmt.Printf("开始执行 %d 个并发Upsert操作...\n", len(docsToUpsert))

    // 3. 启动goroutine并发执行Upsert
    for i, doc := range docsToUpsert {
        wg.Add(1) // 增加WaitGroup计数
        go func(index int, d MyDocument) {
            defer wg.Done() // goroutine完成后减少WaitGroup计数

            // 为每个goroutine克隆一个新的会话
            // 这确保了线程安全，并允许mgo管理底层连接池
            s := session.Clone()
            defer s.Close() // 确保克隆的会话在goroutine结束时关闭

            col := s.DB("testdb").C("mydocs")

            // 定义Upsert的选择器（这里使用Key字段作为唯一标识）
            selector := bson.M{"key": d.Key}

            // 定义更新操作
            // $set 用于更新字段，如果文档存在则更新，不存在则设置
            // $setOnInsert 用于在文档插入时设置字段，如果文档已存在则忽略
            update := bson.M{
                "$set": bson.M{
                    "value":     d.Value,
                    "updatedAt": time.Now(),
                },
                "$setOnInsert": bson.M{
                    "createdAt": time.Now(), // 仅在插入新文档时设置创建时间
                },
            }

            // 执行Upsert操作
            changeInfo, err := col.Upsert(selector, update)
            if err != nil {
                errors <- fmt.Errorf("Upsert操作失败 (文档索引: %d, Key: %s): %v", index, d.Key, err)
                return
            }

            // 打印Upsert结果
            if changeInfo.Matched > 0 {
                fmt.Printf("文档 (Key: %s) 已更新. 匹配: %d, 更新: %d\n", d.Key, changeInfo.Matched, changeInfo.Updated)
            } else if changeInfo.UpsertedId != nil {
                fmt.Printf("新文档 (Key: %s) 已插入，ID: %v\n", d.Key, changeInfo.UpsertedId)
            } else {
                // 理论上不常发生，除非文档匹配但内容没有变化
                fmt.Printf("文档 (Key: %s) 执行Upsert，但无匹配或插入信息\n", d.Key)
            }

        }(i, doc) // 将循环变量和文档作为参数传递给goroutine，避免闭包问题
    }

    // 4. 等待所有goroutine完成
    wg.Wait()
    close(errors) // 关闭错误通道，表示所有错误已发送完毕

    // 5. 收集并打印所有错误
    hasErrors := false
    for err := range errors {
        log.Printf("错误信息: %v", err)
        hasErrors = true
    }

    if !hasErrors {
        fmt.Println("所有Upsert操作已成功完成。")
    }

    // 6. 验证MongoDB中的数据 (可选)
    fmt.Println("\n验证MongoDB中的文档:")
    var results []MyDocument
    err = collection.Find(nil).Sort("key").All(&results) // 按Key排序便于查看
    if err != nil {
        log.Fatalf("检索文档失败: %v", err)
    }
    for _, r := range results {
        fmt.Printf("ID: %v, Key: %s, Value: %s, CreatedAt: %v, UpdatedAt: %v\n",
            r.ID, r.Key, r.Value, r.CreatedAt.Format(time.RFC3339), r.UpdatedAt.Format(time.RFC3339))
    }
}