1
0
mirror of https://github.com/redis/go-redis.git synced 2025-07-29 17:41:15 +03:00

add context cancelation support for blocking operations

This commit is contained in:
Pavlov Aleksey
2020-09-14 21:27:26 +03:00
committed by Vladimir Mihailenco
parent 4b9d6dfd40
commit eda1f9c6ad
2 changed files with 69 additions and 4 deletions

View File

@ -49,7 +49,13 @@ func (hs hooks) process(
ctx context.Context, cmd Cmder, fn func(context.Context, Cmder) error,
) error {
if len(hs.hooks) == 0 {
return fn(ctx, cmd)
return hs.withContext(ctx, func() error {
err := fn(ctx, cmd)
if err != nil {
cmd.SetErr(err)
}
return err
})
}
var hookIndex int
@ -63,7 +69,13 @@ func (hs hooks) process(
}
if retErr == nil {
retErr = fn(ctx, cmd)
retErr = hs.withContext(ctx, func() error {
err := fn(ctx, cmd)
if err != nil {
cmd.SetErr(err)
}
return err
})
}
for hookIndex--; hookIndex >= 0; hookIndex-- {
@ -80,7 +92,13 @@ func (hs hooks) processPipeline(
ctx context.Context, cmds []Cmder, fn func(context.Context, []Cmder) error,
) error {
if len(hs.hooks) == 0 {
return fn(ctx, cmds)
return hs.withContext(ctx, func() error {
err := fn(ctx, cmds)
if err != nil {
setCmdsErr(cmds, err)
}
return err
})
}
var hookIndex int
@ -94,7 +112,13 @@ func (hs hooks) processPipeline(
}
if retErr == nil {
retErr = fn(ctx, cmds)
retErr = hs.withContext(ctx, func() error {
err := fn(ctx, cmds)
if err != nil {
setCmdsErr(cmds, err)
}
return err
})
}
for hookIndex--; hookIndex >= 0; hookIndex-- {
@ -114,6 +138,22 @@ func (hs hooks) processTxPipeline(
return hs.processPipeline(ctx, cmds, fn)
}
func (hs hooks) withContext(ctx context.Context, fn func() error) error {
if ctx.Done() == nil {
return fn()
}
errc := make(chan error, 1)
go func() { errc <- fn() }()
select {
case <-ctx.Done():
return ctx.Err()
case err := <-errc:
return err
}
}
//------------------------------------------------------------------------------
type baseClient struct {