Golang Mgo 高并发写入的性能优化与流量控制

package main

import (
    "fmt"
    "log"
    "time"

    "gopkg.in/mgo.v2" // 注意：这里使用的是gopkg.in/mgo.v2，与labix.org/v2/mgo等效
    "gopkg.in/mgo.v2/bson"
)

// 定义数据结构
type Dog struct {
    Breed string `bson:"breed"`
}

type Person struct {
    Name string    `bson:"name"`
    Pet  Dog       `bson:",inline"`
    Ts   time.Time `bson:"ts"`
}

// insert函数现在接收一个复制的会话
func insert(s *mgo.Session, p Person) {
    // 确保会话在使用完毕后关闭
    defer s.Close()

    err := s.DB("db_log").C("people").Insert(&p)
    if err != nil {
        // 不再panic，而是记录错误，以便主程序可以处理
        log.Printf("Error inserting %s: %v", p.Name, err)
    }
}

func main() {
    // 初始化主会话，只进行一次
    session, err := mgo.Dial("localhost:27017")
    if err != nil {
        log.Fatalf("Could not connect to MongoDB: %v", err)
    }
    defer session.Close() // 确保主会话在程序结束时关闭

    // 设置Mgo连接池的最大连接数，根据MongoDB和应用负载调整
    session.SetPoolLimit(1024) // 示例：设置连接池限制

    // 准备数据
    bob := Person{Name: "Robert", Pet: Dog{Breed: "Labrador"}, Ts: time.Now()}

    for i := 0; i < 100000; i++ { // 模拟大量写入
        // 每次操作前复制会话
        sCopy := session.Copy()
        bob.Name = fmt.Sprintf("Robert-%d", i) // 每次写入不同的数据
        bob.Ts = time.Now()
        go insert(sCopy, bob) // 在新的goroutine中执行插入
        time.Sleep(time.Microsecond * 10) // 稍微放慢生产速度，避免瞬间产生过多goroutine
    }

    // 等待一段时间，让所有goroutine完成
    time.Sleep(time.Second * 5)
    fmt.Println("Finished attempting inserts.")
}

package main

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

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

// 定义数据结构 (与之前相同)
type Dog struct {
    Breed string `bson:"breed"`
}

type Person struct {
    Name string    `bson:"name"`
    Pet  Dog       `bson:",inline"`
    Ts   time.Time `bson:"ts"`
}

// 消费者 goroutine，负责从通道读取数据并写入MongoDB
func worker(id int, session *mgo.Session, dataCh <-chan Person, wg *sync.WaitGroup) {
    defer wg.Done()
    for p := range dataCh {
        sCopy := session.Copy() // 每个写入操作使用独立的会话副本
        err := sCopy.DB("db_log").C("people").Insert(&p)
        sCopy.Close() // 确保会话副本关闭并返回连接池

        if err != nil {
            log.Printf("Worker %d: Error inserting %s: %v", id, p.Name, err)
        } else {
            // fmt.Printf("Worker %d: Successfully inserted %s\n", id, p.Name)
        }
    }
    log.Printf("Worker %d finished.", id)
}

func main() {
    // 设置GOMAXPROCS，充分利用多核CPU
    runtime.GOMAXPROCS(runtime.NumCPU())

    session, err := mgo.Dial("localhost:27017")
    if err != nil {
        log.Fatalf("Could not connect to MongoDB: %v", err)
    }
    defer session.Close()
    session.SetPoolLimit(256) // 根据实际并发写入量调整连接池大小

    // 创建带缓冲的通道，用于传输待写入的数据
    // 缓冲大小决定了生产者在消费者处理数据之前可以“领先”多少个请求
    const channelBufferSize = 1000 // 缓冲区大小
    dataChannel := make(chan Person, channelBufferSize)

    // 使用 WaitGroup 等待所有消费者 goroutine 完成
    var wg sync.WaitGroup

    // 启动多个消费者 goroutine
    const numWorkers = 5 // 并发写入的goroutine数量，根据MongoDB性能和CPU核数调整
    for i := 1; i <= numWorkers; i++ {
        wg.Add(1)
        go worker(i, session, dataChannel, &wg)
    }

    // 生产者 goroutine：生成数据并发送到通道
    const totalInserts = 100000
    fmt.Printf("Starting to produce %d items...\n", totalInserts)
    for i := 0; i < totalInserts; i++ {
        p := Person{
            Name: fmt.Sprintf("Robert-%d", i),
            Pet:  Dog{Breed: "Bulldog"},
            Ts:   time.Now(),
        }
        dataChannel <- p // 当通道满时，此操作会阻塞，实现流量控制
    }
    fmt.Printf("Finished producing %d items.\n", totalInserts)

    // 关闭通道，通知消费者没有更多数据了
    close(dataChannel)

    // 等待所有消费者 goroutine 完成
    wg.Wait()
    fmt.Println("All workers finished and all data processed.")
}

// 示例：配置写入安全级别
sCopy := session.Copy()
defer sCopy.Close()

// 设置写入安全级别为 "fire and forget" (不等待MongoDB确认)
// 极高吞吐量，但有数据丢失风险
sCopy.SetSafe(&mgo.Safe{W: 0})

// 或者设置等待写入到主节点 (默认行为)
// sCopy.SetSafe(&mgo.Safe{W: 1})

// 或者设置等待写入到N个节点，并设置超时时间
// sCopy.SetSafe(&mgo.Safe{W: 2, WTimeout: 5 * time.Second})

// 设置全局的Mgo连接超时时间 (适用于所有操作，不仅仅是写入)
// session.SetSyncTimeout(10 * time.Second) // 同步操作超时
// session.SetSocketTimeout(10 * time.Second) // socket读写超时

err := sCopy.DB("db_log").C("people").Insert(&p)
if err != nil {
    log.Printf("Error inserting with Safe settings: %v", err)
}