add context cancelation support for blocking operations

redis.go

@@ -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 {