Gorm-BUG：PrepareStmt：true 在 SQL 执行时错误 PreparedSQL 会内存泄露

已提交到 issues

github issues

BUG 原因

开启预编译后，gorm (db *PreparedStmtDB) prepare 会拿预编译语句先到 stmts 查找是否存在预编译 (stmt) 对象，有则直接返回 stmt 对象，无则缓存预编译对象，再缓存预编译语句，(db *PreparedStmtDB) ExecContext 有报错则 delete(db.Stmts, query) 删除预编译对象，但未删除缓存预编译语句，会造成 gorm 发送预编译通过但 SQL 执行时错误(如不可重复) PreparedSQL 一直增长未释放数据

版本

• go：1.22
• gorm：1.25.8

复现

1. gorm 配置

   DB, _= gorm.Open(mysql.New(conf), &gorm.Config{
        PrepareStmt: true,
    })

2. 调用创建方法

   database.DB.Create(&m)

3. 报错

   Duplicate entry

Gorm 预编译源码

type Stmt struct {
    *sql.Stmt
    Transaction bool
    prepared    chan struct{}
    prepareErr  error
}

type PreparedStmtDB struct {
    Stmts       map[string]*Stmt
    PreparedSQL []string
    Mux         *sync.RWMutex
    ConnPool
}

func NewPreparedStmtDB(connPool ConnPool) *PreparedStmtDB {
    return &PreparedStmtDB{
        ConnPool:    connPool,
        Stmts:       make(map[string]*Stmt),
        Mux:         &sync.RWMutex{},
        PreparedSQL: make([]string, 0, 100),
    }
}

func (db *PreparedStmtDB) GetDBConn() (*sql.DB, error) {
    if sqldb, ok := db.ConnPool.(*sql.DB); ok {
        return sqldb, nil
    }

    if dbConnector, ok := db.ConnPool.(GetDBConnector); ok && dbConnector != nil {
        return dbConnector.GetDBConn()
    }

    return nil, ErrInvalidDB
}

func (db *PreparedStmtDB) Close() {
    db.Mux.Lock()
    defer db.Mux.Unlock()

    for _, query := range db.PreparedSQL {
        if stmt, ok := db.Stmts[query]; ok {
            delete(db.Stmts, query)
            go stmt.Close()
        }
    }
}

func (sdb *PreparedStmtDB) Reset() {
    sdb.Mux.Lock()
    defer sdb.Mux.Unlock()

    for _, stmt := range sdb.Stmts {
        go stmt.Close()
    }
    sdb.PreparedSQL = make([]string, 0, 100)
    sdb.Stmts = make(map[string]*Stmt)
}

func (db *PreparedStmtDB) prepare(ctx context.Context, conn ConnPool, isTransaction bool, query string) (Stmt, error) {
    db.Mux.RLock()
    if stmt, ok := db.Stmts[query]; ok && (!stmt.Transaction || isTransaction) {
        db.Mux.RUnlock()
        // wait for other goroutines prepared
        <-stmt.prepared
        if stmt.prepareErr != nil {
            return Stmt{}, stmt.prepareErr
        }

        return *stmt, nil
    }
    db.Mux.RUnlock()

    db.Mux.Lock()
    // double check
    if stmt, ok := db.Stmts[query]; ok && (!stmt.Transaction || isTransaction) {
        db.Mux.Unlock()
        // wait for other goroutines prepared
        <-stmt.prepared
        if stmt.prepareErr != nil {
            return Stmt{}, stmt.prepareErr
        }

        return *stmt, nil
    }
    _, ok := db.Stmts[query]

    // cache preparing stmt first
    cacheStmt := Stmt{Transaction: isTransaction, prepared: make(chan struct{})}
    // 设置缓存
    db.Stmts[query] = &cacheStmt
    db.Mux.Unlock()
    _, ok = db.Stmts[query]
    // prepare completed
    defer close(cacheStmt.prepared)

    // Reason why cannot lock conn.PrepareContext
    // suppose the maxopen is 1, g1 is creating record and g2 is querying record.
    // 1. g1 begin tx, g1 is requeue because of waiting for the system call, now `db.ConnPool` db.numOpen == 1.
    // 2. g2 select lock `conn.PrepareContext(ctx, query)`, now db.numOpen == db.maxOpen , wait for release.
    // 3. g1 tx exec insert, wait for unlock `conn.PrepareContext(ctx, query)` to finish tx and release.
    stmt, err := conn.PrepareContext(ctx, query)
    if err != nil {
        cacheStmt.prepareErr = err
        db.Mux.Lock()
        delete(db.Stmts, query)
        db.Mux.Unlock()
        return Stmt{}, err
    }
    db.Mux.Lock()
    cacheStmt.Stmt = stmt
    // 设置缓存 SQL
    db.PreparedSQL = append(db.PreparedSQL, query)
    db.Mux.Unlock()

    return cacheStmt, nil
}

func (db *PreparedStmtDB) ExecContext(ctx context.Context, query string, args ...interface{}) (result sql.Result, err error) {
    stmt, err := db.prepare(ctx, db.ConnPool, false, query)
    if err == nil {
       result, err = stmt.ExecContext(ctx, args...)
       if err != nil {
          db.Mux.Lock()
          defer db.Mux.Unlock()
          go stmt.Close()
          // 删除缓存
          delete(db.Stmts, query)
       }
    }
    return result, err
}

本作品采用《CC 协议》，转载必须注明作者和本文链接