Skip to content

Commit ac4a720

Browse files
noel2004noel2004
committed
refactoring client
1 parent 1944798 commit ac4a720

File tree

5 files changed

+54
-37
lines changed

5 files changed

+54
-37
lines changed

‎coordinator/internal/controller/proxy/client.go‎

Lines changed: 25 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -18,12 +18,25 @@ import (
1818
"scroll-tech/coordinator/internal/types"
1919
)
2020

21+
type ProxyCli interface {
22+
Login(ctx context.Context, genLogin func(string) (*types.LoginParameter, error)) (*ctypes.Response, error)
23+
ProxyLogin(ctx context.Context, param *types.LoginParameter) (*ctypes.Response, error)
24+
Token() string
25+
Reset()
26+
}
27+
28+
type ProverCli interface {
29+
GetTask(ctx context.Context, param *types.GetTaskParameter) (*ctypes.Response, error)
30+
SubmitProof(ctx context.Context, param *types.SubmitProofParameter) (*ctypes.Response, error)
31+
}
32+
2133
// Client wraps an http client with a preset host for coordinator API calls
2234
type upClient struct {
2335
httpClient *http.Client
2436
baseURL string
2537
loginToken string
2638
compatibileMode bool
39+
resetFromMgr func()
2740
}
2841

2942
// NewClient creates a new Client with the specified host
@@ -37,6 +50,12 @@ func newUpClient(cfg *config.UpStream) *upClient {
3750
}
3851
}
3952

53+
func (c *upClient) Reset() {
54+
if c.resetFromMgr != nil {
55+
c.resetFromMgr()
56+
}
57+
}
58+
4059
func (c *upClient) Token() string {
4160
return c.loginToken
4261
}
@@ -175,7 +194,7 @@ func (c *upClient) ProxyLogin(ctx context.Context, param *types.LoginParameter)
175194
}
176195

177196
// GetTask makes a POST request to /v1/get_task with GetTaskParameter
178-
func (c *upClient) GetTask(ctx context.Context, param *types.GetTaskParameter, token string) (*ctypes.Response, error) {
197+
func (c *upClient) GetTask(ctx context.Context, param *types.GetTaskParameter) (*ctypes.Response, error) {
179198
url := fmt.Sprintf("%s/coordinator/v1/get_task", c.baseURL)
180199

181200
jsonData, err := json.Marshal(param)
@@ -189,8 +208,8 @@ func (c *upClient) GetTask(ctx context.Context, param *types.GetTaskParameter, t
189208
}
190209

191210
req.Header.Set("Content-Type", "application/json")
192-
if token != "" {
193-
req.Header.Set("Authorization", "Bearer "+token)
211+
if c.loginToken != "" {
212+
req.Header.Set("Authorization", "Bearer "+c.loginToken)
194213
}
195214

196215
resp, err := c.httpClient.Do(req)
@@ -201,7 +220,7 @@ func (c *upClient) GetTask(ctx context.Context, param *types.GetTaskParameter, t
201220
}
202221

203222
// SubmitProof makes a POST request to /v1/submit_proof with SubmitProofParameter
204-
func (c *upClient) SubmitProof(ctx context.Context, param *types.SubmitProofParameter, token string) (*ctypes.Response, error) {
223+
func (c *upClient) SubmitProof(ctx context.Context, param *types.SubmitProofParameter) (*ctypes.Response, error) {
205224
url := fmt.Sprintf("%s/coordinator/v1/submit_proof", c.baseURL)
206225

207226
jsonData, err := json.Marshal(param)
@@ -215,8 +234,8 @@ func (c *upClient) SubmitProof(ctx context.Context, param *types.SubmitProofPara
215234
}
216235

217236
req.Header.Set("Content-Type", "application/json")
218-
if token != "" {
219-
req.Header.Set("Authorization", "Bearer "+token)
237+
if c.loginToken != "" {
238+
req.Header.Set("Authorization", "Bearer "+c.loginToken)
220239
}
221240

222241
resp, err := c.httpClient.Do(req)

‎coordinator/internal/controller/proxy/client_manager.go‎

Lines changed: 20 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -18,8 +18,11 @@ import (
1818
)
1919

2020
type Client interface {
21-
Client(context.Context) *upClient
22-
Reset(cli *upClient)
21+
// a client to access upstream coordinator with specified identity
22+
// so prover can contact with coordinator as itself
23+
Client(string) ProverCli
24+
// the client to access upstream as proxy itself
25+
ClientAsProxy(context.Context) ProxyCli
2326
Name() string
2427
}
2528

@@ -122,20 +125,17 @@ func (cliMgr *ClientManager) doLogin(ctx context.Context, loginCli *upClient) {
122125
}
123126
}
124127

125-
func (cliMgr *ClientManager) Reset(cli *upClient) {
126-
cliMgr.cachedCli.Lock()
127-
if cliMgr.cachedCli.cli == cli {
128-
log.Info("cached client cleared", "name", cliMgr.name)
129-
cliMgr.cachedCli.cli = nil
130-
}
131-
cliMgr.cachedCli.Unlock()
132-
}
133-
134128
func (cliMgr *ClientManager) Name() string {
135129
return cliMgr.name
136130
}
137131

138-
func (cliMgr *ClientManager) Client(ctx context.Context) *upClient {
132+
func (cliMgr *ClientManager) Client(token string) ProverCli {
133+
loginCli := newUpClient(cliMgr.cfg)
134+
loginCli.loginToken = token
135+
return loginCli
136+
}
137+
138+
func (cliMgr *ClientManager) ClientAsProxy(ctx context.Context) ProxyCli {
139139
cliMgr.cachedCli.RLock()
140140
if cliMgr.cachedCli.cli != nil {
141141
defer cliMgr.cachedCli.RUnlock()
@@ -157,6 +157,14 @@ func (cliMgr *ClientManager) Client(ctx context.Context) *upClient {
157157
// Set new completion context and launch login goroutine
158158
ctx, completionDone := context.WithCancel(context.TODO())
159159
loginCli := newUpClient(cliMgr.cfg)
160+
loginCli.resetFromMgr = func() {
161+
cliMgr.cachedCli.Lock()
162+
if cliMgr.cachedCli.cli == loginCli {
163+
log.Info("cached client cleared", "name", cliMgr.name)
164+
cliMgr.cachedCli.cli = nil
165+
}
166+
cliMgr.cachedCli.Unlock()
167+
}
160168
completionCtx = context.WithValue(ctx, loginCliKey, loginCli)
161169
cliMgr.cachedCli.completionCtx = completionCtx
162170

‎coordinator/internal/controller/proxy/prover_session.go‎

Lines changed: 7 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -149,7 +149,7 @@ func (c *proverSession) maintainLogin(ctx context.Context, cliMgr Client, up str
149149

150150
log.Debug("start proxy login process", "upstream", up, "cli", param.Message.ProverName)
151151

152-
cli := cliMgr.Client(ctx)
152+
cli := cliMgr.ClientAsProxy(ctx)
153153
if cli == nil {
154154
nerr = fmt.Errorf("get upstream cli fail")
155155
return
@@ -163,8 +163,8 @@ func (c *proverSession) maintainLogin(ctx context.Context, cliMgr Client, up str
163163

164164
if resp.ErrCode == ctypes.ErrJWTTokenExpired {
165165
log.Info("up stream has expired, renew upstream connection", "up", up)
166-
cliMgr.Reset(cli)
167-
cli = cliMgr.Client(ctx)
166+
cli.Reset()
167+
cli = cliMgr.ClientAsProxy(ctx)
168168
if cli == nil {
169169
nerr = fmt.Errorf("get upstream cli fail (secondary try)")
170170
return
@@ -226,13 +226,8 @@ func (c *proverSession) GetTask(ctx context.Context, param *types.GetTaskParamet
226226
token := c.proverToken[up]
227227
c.RUnlock()
228228

229-
cli := cliMgr.Client(ctx)
230-
if cli == nil {
231-
return nil, fmt.Errorf("get upstream cli fail")
232-
}
233-
234229
if token.LoginSchema != nil {
235-
resp, err := cli.GetTask(ctx, param, token.Token)
230+
resp, err := cliMgr.Client(token.Token).GetTask(ctx, param)
236231
if err != nil {
237232
return nil, err
238233
}
@@ -253,7 +248,7 @@ func (c *proverSession) GetTask(ctx context.Context, param *types.GetTaskParamet
253248
return nil, fmt.Errorf("update prover token fail: %v", err)
254249
}
255250

256-
return cli.GetTask(ctx, param, newToken.Token)
251+
return cliMgr.Client(newToken.Token).GetTask(ctx, param)
257252

258253
}
259254

@@ -264,13 +259,8 @@ func (c *proverSession) SubmitProof(ctx context.Context, param *types.SubmitProo
264259
token := c.proverToken[up]
265260
c.RUnlock()
266261

267-
cli := cliMgr.Client(ctx)
268-
if cli == nil {
269-
return nil, fmt.Errorf("get upstream cli fail")
270-
}
271-
272262
if token.LoginSchema != nil {
273-
resp, err := cli.SubmitProof(ctx, param, token.Token)
263+
resp, err := cliMgr.Client(token.Token).SubmitProof(ctx, param)
274264
if err != nil {
275265
return nil, err
276266
}
@@ -291,5 +281,5 @@ func (c *proverSession) SubmitProof(ctx context.Context, param *types.SubmitProo
291281
return nil, fmt.Errorf("update prover token fail: %v", err)
292282
}
293283

294-
return cli.SubmitProof(ctx, param, newToken.Token)
284+
return cliMgr.Client(newToken.Token).SubmitProof(ctx, param)
295285
}

‎coordinator/internal/orm/orm_test.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,8 +28,8 @@ func TestMain(m *testing.M) {
2828
defer func() {
2929
if testApps != nil {
3030
testApps.Free()
31+
tearDownEnv(t)
3132
}
32-
tearDownEnv(t)
3333
}()
3434
m.Run()
3535
}

‎coordinator/test/proxy_test.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -63,7 +63,7 @@ func testProxyClient(t *testing.T) {
6363
defer cancel()
6464

6565
// Test Client method
66-
client := clientManager.Client(ctx)
66+
client := clientManager.ClientAsProxy(ctx)
6767

6868
// Client should not be nil if login succeeds
6969
// Note: This might be nil if the coordinator is not properly set up for proxy authentication

0 commit comments

Comments
 (0)