Skip to content

Commit 9df6429

Browse files
noel2004noel2004
committed
wip
1 parent e6be62f commit 9df6429

File tree

1 file changed

+96
-70
lines changed

1 file changed

+96
-70
lines changed

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

Lines changed: 96 additions & 70 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,8 @@
11
package proxy
22

33
import (
4-
"bytes"
54
"context"
6-
"encoding/json"
75
"fmt"
8-
"net/http"
96
"sync"
107

118
ctypes "scroll-tech/common/types"
@@ -15,126 +12,155 @@ import (
1512
// Client wraps an http client with a preset host for coordinator API calls
1613
type proverSession struct {
1714
sync.RWMutex
18-
proverToken string
15+
proverToken string
16+
phase uint
17+
completionCtx context.Context
1918
}
2019

21-
func (c *proverSession) doProverLogin(ctx context.Context, cliMgr Client, param *types.LoginParameter) (*types.LoginSchema, error) {
20+
func (c *proverSession) maintainLogin(ctx context.Context, cliMgr Client, param *types.LoginParameter, phase uint) error {
21+
c.Lock()
22+
curPhase := c.phase
23+
if c.completionCtx != nil {
24+
waitctx := c.completionCtx
25+
c.Unlock()
26+
select {
27+
case <-waitctx.Done():
28+
return c.maintainLogin(ctx, cliMgr, param, phase)
29+
case <-ctx.Done():
30+
return fmt.Errorf("ctx fail")
31+
}
32+
}
33+
34+
if phase < curPhase {
35+
// outdate login phase, give up
36+
c.Unlock()
37+
return nil
38+
}
39+
40+
// occupy the update slot
41+
completeCtx, cf := context.WithCancel(ctx)
42+
defer cf()
43+
c.completionCtx = completeCtx
44+
c.Unlock()
45+
2246
cli := cliMgr.Client(ctx)
2347
if cli == nil {
24-
return nil, fmt.Errorf("get upstream cli fail")
48+
return fmt.Errorf("get upstream cli fail")
2549
}
2650

27-
// like SDK, we would try one more time if the upstream token is expired
2851
resp, err := cli.ProxyLogin(ctx, param)
2952
if err != nil {
30-
return nil, fmt.Errorf("proxylogin fail: %v", err)
53+
return fmt.Errorf("proxylogin fail: %v", err)
3154
}
3255

3356
if resp.ErrCode == ctypes.ErrJWTTokenExpired {
3457
cliMgr.Reset(cli)
3558
cli = cliMgr.Client(ctx)
3659
if cli == nil {
37-
return nil, fmt.Errorf("get upstream cli fail (secondary try)")
60+
return fmt.Errorf("get upstream cli fail (secondary try)")
3861
}
3962

4063
// like SDK, we would try one more time if the upstream token is expired
4164
resp, err = cli.ProxyLogin(ctx, param)
4265
if err != nil {
43-
return nil, fmt.Errorf("proxylogin fail: %v", err)
66+
return fmt.Errorf("proxylogin fail: %v", err)
4467
}
4568
}
4669

4770
if resp.ErrCode != 0 {
48-
return nil, fmt.Errorf("upstream fail: %d (%s)", resp.ErrCode, resp.ErrMsg)
71+
return fmt.Errorf("upstream fail: %d (%s)", resp.ErrCode, resp.ErrMsg)
4972
}
5073

5174
var loginResult types.LoginSchema
5275
if err := resp.DecodeData(&loginResult); err != nil {
53-
return nil, err
76+
return err
5477
}
5578

56-
return &loginResult, nil
79+
c.Lock()
80+
defer c.Unlock()
81+
c.proverToken = loginResult.Token
82+
c.completionCtx = nil
83+
84+
return nil
5785
}
5886

5987
// ProxyLogin makes a POST request to /v1/proxy_login with LoginParameter
60-
func (c *proverSession) ProxyLogin(ctx context.Context, cli Client, param *types.LoginParameter) (*types.LoginSchema, error) {
61-
url := fmt.Sprintf("%s/coordinator/v1/proxy_login", c.baseURL)
62-
63-
jsonData, err := json.Marshal(param)
64-
if err != nil {
65-
return nil, fmt.Errorf("failed to marshal proxy login parameter: %w", err)
66-
}
88+
func (c *proverSession) ProxyLogin(ctx context.Context, cli Client, param *types.LoginParameter) error {
89+
c.RLock()
90+
phase := c.phase + 1
91+
c.RUnlock()
6792

68-
req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(jsonData))
69-
if err != nil {
70-
return nil, fmt.Errorf("failed to create proxy login request: %w", err)
71-
}
72-
73-
req.Header.Set("Content-Type", "application/json")
74-
req.Header.Set("Authorization", "Bearer "+c.loginToken)
75-
76-
proxyLoginResp, err := c.httpClient.Do(req)
77-
if err != nil {
78-
return nil, fmt.Errorf("failed to perform proxy login request: %w", err)
79-
}
80-
defer proxyLoginResp.Body.Close()
81-
82-
// Call helper's OnResp method with the response
83-
c.helper.OnResp(c, proxyLoginResp)
84-
85-
// Parse proxy login response as LoginSchema
86-
if proxyLoginResp.StatusCode == http.StatusOK {
87-
var loginResult types.LoginSchema
88-
if err := json.NewDecoder(proxyLoginResp.Body).Decode(&loginResult); err == nil {
89-
return &loginResult, nil
90-
}
91-
// If parsing fails, still return success but with nil result
92-
return nil, nil
93-
}
94-
95-
return nil, fmt.Errorf("proxy login request failed with status: %d", proxyLoginResp.StatusCode)
93+
return c.maintainLogin(ctx, cli, param, phase)
9694
}
9795

9896
// GetTask makes a POST request to /v1/get_task with GetTaskParameter
99-
func (c *proverSession) GetTask(ctx context.Context, param *types.GetTaskParameter, token string) (*http.Response, error) {
100-
url := fmt.Sprintf("%s/coordinator/v1/get_task", c.baseURL)
97+
func (c *proverSession) GetTask(ctx context.Context, param *types.GetTaskParameter, cliMgr Client) (*ctypes.Response, error) {
98+
c.RLock()
99+
phase := c.phase
100+
token := c.proverToken
101+
c.RUnlock()
101102

102-
jsonData, err := json.Marshal(param)
103-
if err != nil {
104-
return nil, fmt.Errorf("failed to marshal get task parameter: %w", err)
103+
cli := cliMgr.Client(ctx)
104+
if cli == nil {
105+
return nil, fmt.Errorf("get upstream cli fail")
105106
}
106107

107-
req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(jsonData))
108+
resp, err := cli.GetTask(ctx, param, token)
108109
if err != nil {
109-
return nil, fmt.Errorf("failed to create get task request: %w", err)
110+
return nil, err
110111
}
111112

112-
req.Header.Set("Content-Type", "application/json")
113-
if token != "" {
114-
req.Header.Set("Authorization", "Bearer "+token)
113+
if resp.ErrCode == ctypes.ErrJWTTokenExpired {
114+
// get param from ctx
115+
loginParam, ok := ctx.Value(LoginParamCache).(*types.LoginParameter)
116+
if !ok {
117+
return nil, fmt.Errorf("Unexpected error, no loginparam ctx value")
118+
}
119+
120+
err = c.maintainLogin(ctx, cliMgr, loginParam, phase)
121+
if err != nil {
122+
return nil, fmt.Errorf("update prover token fail: %V", err)
123+
}
124+
125+
// like SDK, we would try one more time if the upstream token is expired
126+
return cli.GetTask(ctx, param, token)
115127
}
116128

117-
return c.httpClient.Do(req)
129+
return resp, nil
118130
}
119131

120132
// SubmitProof makes a POST request to /v1/submit_proof with SubmitProofParameter
121-
func (c *proverSession) SubmitProof(ctx context.Context, param *types.SubmitProofParameter, token string) (*http.Response, error) {
122-
url := fmt.Sprintf("%s/coordinator/v1/submit_proof", c.baseURL)
133+
func (c *proverSession) SubmitProof(ctx context.Context, param *types.SubmitProofParameter, cliMgr Client) (*ctypes.Response, error) {
134+
c.RLock()
135+
phase := c.phase
136+
token := c.proverToken
137+
c.RUnlock()
123138

124-
jsonData, err := json.Marshal(param)
125-
if err != nil {
126-
return nil, fmt.Errorf("failed to marshal submit proof parameter: %w", err)
139+
cli := cliMgr.Client(ctx)
140+
if cli == nil {
141+
return nil, fmt.Errorf("get upstream cli fail")
127142
}
128143

129-
req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(jsonData))
144+
resp, err := cli.SubmitProof(ctx, param, token)
130145
if err != nil {
131-
return nil, fmt.Errorf("failed to create submit proof request: %w", err)
146+
return nil, err
132147
}
133148

134-
req.Header.Set("Content-Type", "application/json")
135-
if token != "" {
136-
req.Header.Set("Authorization", "Bearer "+token)
149+
if resp.ErrCode == ctypes.ErrJWTTokenExpired {
150+
// get param from ctx
151+
loginParam, ok := ctx.Value(LoginParamCache).(*types.LoginParameter)
152+
if !ok {
153+
return nil, fmt.Errorf("Unexpected error, no loginparam ctx value")
154+
}
155+
156+
err = c.maintainLogin(ctx, cliMgr, loginParam, phase)
157+
if err != nil {
158+
return nil, fmt.Errorf("update prover token fail: %V", err)
159+
}
160+
161+
// like SDK, we would try one more time if the upstream token is expired
162+
return cli.SubmitProof(ctx, param, token)
137163
}
138164

139-
return c.httpClient.Do(req)
165+
return resp, nil
140166
}

0 commit comments

Comments
 (0)