context.WithTimeout()之实现Gorm超时控制

心的步伐

已于 2023-09-02 17:48:41 修改

阅读量2.6k

点赞数 2

分类专栏： Go MySQL 文章标签： golang gorm context golang超时控制

于 2023-09-02 17:46:39 首次发布

本文链接：https://blog.csdn.net/cao1315020626/article/details/132640824

版权

Go 同时被 2 个专栏收录

15 篇文章 1 订阅

订阅专栏

MySQL

3 篇文章 2 订阅

订阅专栏

文章目录

1. 写在前面
2. 如何实现超时控制
3. 实现Gorm的超时控制
4. 看向更底层
5. 实现超时控制逻辑
6. Gorm内部的context done校验逻辑
7. 小结
8. 参考文档

1. 写在前面

Context是Golang中的上下文，Gorm是当前用的比较多的SQL组件库，在Gorm中，Gorm通过WithContext提供了Context支持，在某种程度上来说，Gorm目前提供的Context可以实现对于SQL的超时控制。

本意是为了实现gorm单条SQL语句执行时间的控制而引入的context.WithTimeout()，限制单条语句执行时间，后面在修改过程中遇到了关于context.WithTimeout()的问题，在修复问题的过程中让自己对context.WithTimeout()有了更深的理解。

2. 如何实现超时控制

在Gorm中如果将超时控制作用于单条SQL的查询，给出的示例如下：

ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
defer cancel()

db.WithContext(ctx).Find(&users)

对于这个case，采用context.WithTimeout给当前的context加上了2s的超时控制，即在2s内如果这条查询的SQL没有执行完，会得到context deadline exceeded的error。

另外关于db操作取消context的逻辑，golang也给出了相关的示例，具体可以参考：Canceling in-progress operations，即意味着我们再需要进行超时控制的地方调用QueryWithTimeout方法可以实现对查询语句的超时控制。

func QueryWithTimeout(ctx context.Context) {
    // Create a Context with a timeout.
    queryCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
    defer cancel()

    // Pass the timeout Context with a query.
    rows, err := db.QueryContext(queryCtx, "SELECT * FROM album")
    if err != nil {
        log.Fatal(err)
    }
    defer rows.Close()

    // Handle returned rows.
}

3. 实现Gorm的超时控制

Gorm官方文档给出的例子，给出的例子是查询用户信息，而如果我们有多条语句，多个表，按照示例，我们需要写多个ontext.WithTimeout的代码片段来实现。如果在需求的迭代中，某一天组内来了新人，没有加上这个代码片段，则相关的SQL执行就失去了超时控制。这都是因为这个控制是发生在过于上层的地方了。

我们先思考一个问题：如何把超时控制放在更加底层的逻辑，即更接近sql执行逻辑的前置步骤？

Go官方给出的context.WithTimeout的用法时，给出的注释如下：

// WithTimeout returns WithDeadline(parent, time.Now().Add(timeout)).
//
// Canceling this context releases resources associated with it, so code should
// call cancel as soon as the operations running in this Context complete:
// 取消此上下文并释放与它相关的资源，因此在此上下文上下文中运行的操作完成后，代码应立即取消调用
//
// 	func slowOperationWithTimeout(ctx context.Context) (Result, error) {
// 		ctx, cancel := context.WithTimeout(ctx, 100*time.Millisecond)
// 		defer cancel()  // releases resources if slowOperation completes before timeout elapses
// 		return slowOperation(ctx)
// 	}

从这个注释中可以知道，在使用context.WithTimeout的时候，defer cancel()需要尽可能和代码片段写在一起。

从这个层面，会发现，如果把超时控制放在太上层的db方法调用块这，会导致调用context.WithTimeout()的代码过多，意味着重复的代码片段也就越多，能够复用的代码片段就越来越少。

所以，我们需要想办法把这个粒度放到更细层面，最好是能够让所有的db方法调用，最后都可以经过这个方法，然后我们在这个方法进行超时控制逻辑的控制，就可以解决这个问题。

4. 看向更底层

为了实现能够让所有的db方法调用都可以经过这个方法，就意味着这个方法还在更加底层的位置，所以我们继续向下看，

// Find finds all records matching given conditions conds
func (db *DB) Find(dest interface{}, conds ...interface{}) (tx *DB) {
	tx = db.getInstance()
	if len(conds) > 0 {
		if exprs := tx.Statement.BuildCondition(conds[0], conds[1:]...); len(exprs) > 0 {
			tx.Statement.AddClause(clause.Where{Exprs: exprs})
		}
	}
	tx.Statement.Dest = dest
	return tx.callbacks.Query().Execute(tx) // 实际执行的事callbacks中注册的Query的方法
}

func (cs *callbacks) Query() *processor {
	return cs.processors["query"]
}

func RegisterDefaultCallbacks(db *gorm.DB, config *Config) {
	enableTransaction := func(db *gorm.DB) bool {
		return !db.SkipDefaultTransaction
	}

	// ...省略一些代码

	queryCallback := db.Callback().Query()
	queryCallback.Register("gorm:query", Query)
	queryCallback.Register("gorm:preload", Preload)
	queryCallback.Register("gorm:after_query", AfterQuery)
	queryCallback.Clauses = config.QueryClauses

	// ...省略一些代码
}

func (p *processor) Execute(db *DB) *DB {
	// ...
	for _, f := range p.fns {
		f(db)
	}

	//...
	return db
}

我们会发现在上面的代码中，使用到了db.Callback，我们继续看源码会发现，最终所注册的callback方法都会变为p.fns中的函数，而在执行过程中，这些callback方法都会在这里被调用，看到了这里，当即我决定把我们的逻辑放在callback中，在所有callback执行前，调用context.WithTimeout()，而在callback最后一个方法，我们执行cancelFunc，即可实现整个流程的超时控制了。

5. 实现超时控制逻辑

func registerCallbacks(db *gorm.DB) {
	query := db.Callback().Query()
	query.Before("gorm:query").Register("custom:before_gorm_query", BeforeQuery)
	query.After("gorm:after_query").Register("custom:after_gorm_query", AfterQuery)
}

func BeforeQuery(db *gorm.DB) {
	var cancelFunc context.CancelFunc
	db.Statement.Context, cancelFunc = context.WithTimeout(db.Statement.Context, 3*time.Second)
	db.Statement.Context = context.WithValue( // 这里负责把超时控制注入进去
		db.Statement.Context,
		"ctxInfo",
		&CtxInfo{
			cancelFunc: &cancelFunc,
		},
	)
}

func AfterQuery(db *gorm.DB) {
	ctxInfo := getCtxInfo(db.Statement.Context)
	if ctxInfo != nil {
		f := *ctxInfo.cancelFunc
		f() // execute the cancel func to cancel the context
	}
}

// getCtxInfo get ctx info from context
func getCtxInfo(ctx context.Context) *CtxInfo {
	if ctxInfo, ok := ctx.Value("ctxInfo").(*CtxInfo); ok {
		return ctxInfo
	}
	return nil
}

// CtxInfo
type CtxInfo struct {
	cancelFunc *context.CancelFunc
}

上面就是在callback中注入我们想要执行的超时控制逻辑，然后在结束的时候，执行cancelFunc，看起来似乎没问题，我们正常查询db，返回的结果也正常。

但问题来了，当我们复用gorm session执行链式操作(Gorm链式操作)的时候，多条的语句查询却出现了问题，本意是想复用session进行sql查询，但却引入了错误？

例子如下：

o := orm.GetORM().NewGormSession(ctx).Table("user_tab")
// 这一条结果正常
result := o.Where("user_id = ?", req.UserId).Find(user)
// 这一条结果异常，error为context canceled
result = o.Where("user_id = ?", req.UserId).Find(user)
// {"level":"error","ts":"2023-09-02 16:11:31",
// "caller":"user/api_get_user_data_by_id.go:30",
// "message":"ApiGetUserDataById|get user error, err=context canceled"}

仔细看了一下上面的逻辑，发现一个问题，在我们第二次查询发生的时候，它的context已经被取消了，这是为什么呢？这个就和context.WithTimeout()函数本身的实现有关系了，于是我们来看看这个方法的实现：

func WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc) {
	return WithDeadline(parent, time.Now().Add(timeout))
}

func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {
	if parent == nil {
		panic("cannot create context from nil parent")
	}
	if cur, ok := parent.Deadline(); ok && cur.Before(d) {
		// The current deadline is already sooner than the new one.
		return WithCancel(parent) // 当父context有过期时间且过期时间在自己之前，就复用父的
	}
	c := &timerCtx{
		cancelCtx: newCancelCtx(parent),
		deadline:  d,
	}
	propagateCancel(parent, c)
	dur := time.Until(d)
	if dur <= 0 {
		c.cancel(true, DeadlineExceeded, nil) // deadline has already passed
		return c, func() { c.cancel(false, Canceled, nil) }
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.err == nil {
		c.timer = time.AfterFunc(dur, func() {
			c.cancel(true, DeadlineExceeded, nil)
		})
	}
	return c, func() { c.cancel(true, Canceled, nil) }
}

func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
	c := withCancel(parent) // 创建以父类为对象那个的cancelCtx
	return c, func() { c.cancel(true, Canceled, nil) }
}

func withCancel(parent Context) *cancelCtx {
	if parent == nil {
		panic("cannot create context from nil parent")
	}
	c := newCancelCtx(parent)
	propagateCancel(parent, c) // 这里会去查看一下父类是否done了
	return c
}

func propagateCancel(parent Context, child canceler) {
	done := parent.Done() // 这里获取是否done了
	if done == nil {
		return // parent is never canceled
	}

	select {
	case <-done:
		// parent is already canceled // 这里cancel以后，context会直接done了，error是context canceled
		child.cancel(false, parent.Err(), Cause(parent))
		return
	default:
	}
	// 省略......
}

仔细检查发现第二个的context因为db.Statement.Context的缘故，导致复用了第一个的timerCtx和cancelCtx，导致在第二次进行context.WithTimeout()调用的时候，就已经被canceled掉了。这里的原因是因为第一次的查询结束，把自己的cancel掉了，第二次查询把第一次查询的context当做父context，而golang context设计的规则就是如果父context被cancel掉了，子也会被cancel掉，所以导致问题的出现。

问题：context复用导致子context因为父context被cancel而cancel，返回错误context canceled。

解决问题：规避复用，存储最原始的context，后续每一次都用最原始context，从而实现多次查询的超时控制。

func BeforeQuery(db *gorm.DB) {
	var ctxInfo *CtxInfo
	ctxInfo = getCtxInfo(db.Statement.Context)
	if ctxInfo == nil {
		ctxInfo = &CtxInfo{OriginContext: db.Statement.Context}
	}
	var cancelFunc context.CancelFunc
	db.Statement.Context, cancelFunc = context.WithTimeout(ctxInfo.OriginContext, 3*time.Second)
	ctxInfo.cancelFunc = &cancelFunc
	db.Statement.Context = context.WithValue(
		db.Statement.Context,
		"ctxInfo",
		ctxInfo,
	)
}

// CtxInfo
type CtxInfo struct {
	cancelFunc    *context.CancelFunc
	OriginContext context.Context
}

存储最原始的context，然后后续每一次查询都使用最原始的context即可解决这个问题，且保证了每一个sql都有超时控制。

6. Gorm内部的context done校验逻辑

我们对context进行了超时控制，gorm如何识别到这个context完成了或者取消了呢？

在Gorm的源码中我们主要关注这几个方法就可以知道gorm是如何实现的了。

func (c *connector) Connect(ctx context.Context) (driver.Conn, error) {
	// ...省略
	// Call startWatcher for context support (From Go 1.8)
	mc.startWatcher()
	// ...省略
	
	return mc, nil
}

func (mc *mysqlConn) startWatcher() {
	watcher := make(chan context.Context, 1)
	mc.watcher = watcher
	finished := make(chan struct{})
	mc.finished = finished
	go func() {
		for {
			var ctx context.Context
			select {
			case ctx = <-watcher: // 这里的watcher，在后续的执行sql过程中，会将ctx塞入
			case <-mc.closech:
				return
			}

			select {
			case <-ctx.Done():
				mc.cancel(ctx.Err())
			case <-finished:
			case <-mc.closech:
				return
			}
		}
	}()
}

func (mc *mysqlConn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Rows, error) {
	dargs, err := namedValueToValue(args)
	if err != nil {
		return nil, err
	}

	if err := mc.watchCancel(ctx); err != nil {// 这里watchCancel
		return nil, err
	}

	rows, err := mc.query(query, dargs)
	if err != nil {
		mc.finish()
		return nil, err
	}
	rows.finish = mc.finish
	return rows, err
}

func (mc *mysqlConn) watchCancel(ctx context.Context) error {
	if mc.watching {
		// Reach here if canceled,
		// so the connection is already invalid
		mc.cleanup()
		return nil
	}
	// When ctx is already cancelled, don't watch it.
	if err := ctx.Err(); err != nil {
		return err
	}
	// When ctx is not cancellable, don't watch it.
	if ctx.Done() == nil {
		return nil
	}
	// When watcher is not alive, can't watch it.
	if mc.watcher == nil {
		return nil
	}

	mc.watching = true
	mc.watcher <- ctx// 这里塞入
	return nil
}

上面的逻辑帮助在connection层面监测context的变化，gorm在过程的处理中也一直在重复这下面这段代码的操作，在很多地方都会判断context是否已经结束操作，如果是的话，会直接返回。

select {
default:
case <-ctx.Done():
	return nil, ctx.Err()
}

上述两部分的逻辑，实现的gorm的context控制逻辑。

7. 小结

对Gorm执行sql加超时控制，最开始以为可以通过实现ConnPool的四个方法加入超时控制逻辑来实现，后面发现ConnPool的方法也是sql执行生命周期中的一环，想要控制超时时间还是需要从callback入手。

后续加入了超时控制，对于单条sql的查询没有问题，因为context只用一次，不会引起context的嵌套问题，但忽略Gorm复用session的环节，导致在多次查询的情况下，除第一次外，后续的context由于嵌套且父context被手动cancel了，在执行sql的时候，由于context已经done了，导致查询失败。

在一开始看到这个问题，因为失败的非常快，就想着肯定有提前取消的问题，刚开始是将AfterQuery中的cancelFunc的执行逻辑去除了，发现这样就没有问题。于是上手debug，发现cancel函数只能释放相关资源(定时器资源)，但并不会将自己从context链中移除，导致第二次执行WithTimeout的时候，由于context的复用，导致context对象出现了两份timerCtx和cancelCtx，就是就发现是context的复用问题。

加上对context的理解，cancel函数如果父类有，往下传递过程中，子类如果也有，且父类的过期时间早于子类的话，会在创建cancelCtx的时候，提前判断父是否Done了，如果Done了，此时子会立即取消自己，然后返回。

后面加上了暂存最原始context的逻辑之后，问题就解决了。