From 662c654fe3d2cc90e3bdf2d4d879de1596cf9b2b Mon Sep 17 00:00:00 2001
From: SimFG <bang.fu@zilliz.com>
Date: Sat, 24 Sep 2022 09:56:51 +0800
Subject: [PATCH] Add the `context` param for the `hook` interceptor (#19405)

Signed-off-by: SimFG <bang.fu@zilliz.com>

Signed-off-by: SimFG <bang.fu@zilliz.com>
---
 api/hook/hook.go                        |  8 +++++---
 internal/proxy/hook_interceptor.go      | 12 ++++++------
 internal/proxy/hook_interceptor_test.go |  6 +++---
 3 files changed, 14 insertions(+), 12 deletions(-)

diff --git a/api/hook/hook.go b/api/hook/hook.go
index 45254385b42f7..c606424564559 100644
--- a/api/hook/hook.go
+++ b/api/hook/hook.go
@@ -1,9 +1,11 @@
 package hook
 
+import "context"
+
 type Hook interface {
 	Init(params map[string]string) error
-	Mock(req interface{}, fullMethod string) (bool, interface{}, error)
-	Before(req interface{}, fullMethod string) error
-	After(result interface{}, err error, fullMethod string) error
+	Mock(ctx context.Context, req interface{}, fullMethod string) (bool, interface{}, error)
+	Before(ctx context.Context, req interface{}, fullMethod string) error
+	After(ctx context.Context, result interface{}, err error, fullMethod string) error
 	Release()
 }
diff --git a/internal/proxy/hook_interceptor.go b/internal/proxy/hook_interceptor.go
index 3e548f076c72d..ba24ade66e366 100644
--- a/internal/proxy/hook_interceptor.go
+++ b/internal/proxy/hook_interceptor.go
@@ -18,15 +18,15 @@ func (d defaultHook) Init(params map[string]string) error {
 	return nil
 }
 
-func (d defaultHook) Mock(req interface{}, fullMethod string) (bool, interface{}, error) {
+func (d defaultHook) Mock(ctx context.Context, req interface{}, fullMethod string) (bool, interface{}, error) {
 	return false, nil, nil
 }
 
-func (d defaultHook) Before(req interface{}, fullMethod string) error {
+func (d defaultHook) Before(ctx context.Context, req interface{}, fullMethod string) error {
 	return nil
 }
 
-func (d defaultHook) After(result interface{}, err error, fullMethod string) error {
+func (d defaultHook) After(ctx context.Context, result interface{}, err error, fullMethod string) error {
 	return nil
 }
 
@@ -79,15 +79,15 @@ func UnaryServerHookInterceptor() grpc.UnaryServerInterceptor {
 			err        error
 		)
 
-		if isMock, mockResp, err = hoo.Mock(req, fullMethod); isMock {
+		if isMock, mockResp, err = hoo.Mock(ctx, req, fullMethod); isMock {
 			return mockResp, err
 		}
 
-		if err = hoo.Before(req, fullMethod); err != nil {
+		if err = hoo.Before(ctx, req, fullMethod); err != nil {
 			return nil, err
 		}
 		realResp, realErr = handler(ctx, req)
-		if err = hoo.After(realResp, realErr, fullMethod); err != nil {
+		if err = hoo.After(ctx, realResp, realErr, fullMethod); err != nil {
 			return nil, err
 		}
 		return realResp, realErr
diff --git a/internal/proxy/hook_interceptor_test.go b/internal/proxy/hook_interceptor_test.go
index bf06aac4b9240..7ea1acf667148 100644
--- a/internal/proxy/hook_interceptor_test.go
+++ b/internal/proxy/hook_interceptor_test.go
@@ -28,7 +28,7 @@ type mockHook struct {
 	mockErr error
 }
 
-func (m mockHook) Mock(req interface{}, fullMethod string) (bool, interface{}, error) {
+func (m mockHook) Mock(ctx context.Context, req interface{}, fullMethod string) (bool, interface{}, error) {
 	return true, m.mockRes, m.mockErr
 }
 
@@ -42,7 +42,7 @@ type beforeMock struct {
 	err    error
 }
 
-func (b beforeMock) Before(r interface{}, fullMethod string) error {
+func (b beforeMock) Before(ctx context.Context, r interface{}, fullMethod string) error {
 	re, ok := r.(*req)
 	if !ok {
 		return errors.New("r is invalid type")
@@ -61,7 +61,7 @@ type afterMock struct {
 	err    error
 }
 
-func (a afterMock) After(r interface{}, err error, fullMethod string) error {
+func (a afterMock) After(ctx context.Context, r interface{}, err error, fullMethod string) error {
 	re, ok := r.(*resp)
 	if !ok {
 		return errors.New("r is invalid type")