tikv
/
client-go
mirror of https://github.com/tikv/client-go.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

client: support interceptor in SendRequestAsync (#1735) 

Signed-off-by: zyguan <zhongyangguan@gmail.com>
This commit is contained in:
zyguan 2025-08-22 15:58:38 +08:00 committed by GitHub
parent a05a5382fb
commit c0b6188ffd
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 8 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
internal/client/client_interceptor.go
View File

@ -19,7 +19,6 @@ import (
"sync/atomic" "sync/atomic"
"time" "time"
"github.com/tikv/client-go/v2/internal/logutil"
"github.com/tikv/client-go/v2/internal/resourcecontrol" "github.com/tikv/client-go/v2/internal/resourcecontrol"
"github.com/tikv/client-go/v2/tikvrpc" "github.com/tikv/client-go/v2/tikvrpc"
"github.com/tikv/client-go/v2/tikvrpc/interceptor" "github.com/tikv/client-go/v2/tikvrpc/interceptor"
@ -114,6 +113,14 @@ func (r interceptedClient) SendRequestAsync(ctx context.Context, addr string, re
} }
cb.Inject(func(resp *tikvrpc.Response, err error) (*tikvrpc.Response, error) { cb.Inject(func(resp *tikvrpc.Response, err error) (*tikvrpc.Response, error) {
if ctxInterceptor := interceptor.GetRPCInterceptorFromCtx(ctx); ctxInterceptor != nil {
// TODO(zyguan): In async API, the interceptor is only triggered upon receiving the response. Maybe
// support AsyncRPCInterceptor later.
getResp := func(target string, req *tikvrpc.Request) (*tikvrpc.Response, error) {
return resp, err
}
resp, err = ctxInterceptor.Wrap(getResp)(addr, req)
}
if resp != nil { if resp != nil {
respInfo := resourcecontrol.MakeResponseInfo(resp) respInfo := resourcecontrol.MakeResponseInfo(resp)
consumption, waitDuration, err := resourceControlInterceptor.OnResponseWait(ctx, resourceGroupName, reqInfo, respInfo) consumption, waitDuration, err := resourceControlInterceptor.OnResponseWait(ctx, resourceGroupName, reqInfo, respInfo)
@ -128,11 +135,6 @@ func (r interceptedClient) SendRequestAsync(ctx context.Context, addr string, re
}) })
} }
if ctxInterceptor := interceptor.GetRPCInterceptorFromCtx(ctx); ctxInterceptor != nil {
// TODO(zyguan): support AsyncRPCInterceptor
logutil.Logger(ctx).Warn("SendRequestAsync with interceptor is unsupported")
}
r.Client.SendRequestAsync(ctx, addr, req, cb) r.Client.SendRequestAsync(ctx, addr, req, cb)
} }