Skip to content

Commit 8f9f505

Browse files
shizhMSFTshizhMSFT
authored
fix(DynamicStore): retry setCredsStore on next PUT (#728)
Fix #727 Signed-off-by: Shiwei Zhang <[email protected]>
1 parent d3ff5dc commit 8f9f505

3 files changed

Lines changed: 135 additions & 8 deletions

File tree

‎internal/syncutil/once.go‎

Lines changed: 34 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,10 +15,14 @@ limitations under the License.
1515

1616
package syncutil
1717

18-
import "context"
18+
import (
19+
"context"
20+
"sync"
21+
"sync/atomic"
22+
)
1923

2024
// Once is an object that will perform exactly one action.
21-
// Unlike sync.Once, this Once allowes the action to have return values.
25+
// Unlike sync.Once, this Once allows the action to have return values.
2226
type Once struct {
2327
result interface{}
2428
err error
@@ -68,3 +72,31 @@ func (o *Once) Do(ctx context.Context, f func() (interface{}, error)) (bool, int
6872
}
6973
}
7074
}
75+
76+
// OnceOrRetry is an object that will perform exactly one success action.
77+
type OnceOrRetry struct {
78+
done atomic.Bool
79+
lock sync.Mutex
80+
}
81+
82+
// OnceOrRetry calls the function f if and only if Do is being called for the
83+
// first time for this instance of Once or all previous calls to Do are failed.
84+
func (o *OnceOrRetry) Do(f func() error) error {
85+
// fast path
86+
if o.done.Load() {
87+
return nil
88+
}
89+
90+
// slow path
91+
o.lock.Lock()
92+
defer o.lock.Unlock()
93+
94+
if o.done.Load() {
95+
return nil
96+
}
97+
if err := f(); err != nil {
98+
return err
99+
}
100+
o.done.Store(true)
101+
return nil
102+
}

‎internal/syncutil/once_test.go‎

Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ import (
2222
"reflect"
2323
"strconv"
2424
"sync"
25+
"sync/atomic"
2526
"testing"
2627
"time"
2728
)
@@ -191,3 +192,97 @@ func TestOnce_Do_Cancel_Panic(t *testing.T) {
191192
t.Fatalf("Once.Do() result = %v, want %v", result, wantResult)
192193
}
193194
}
195+
196+
func TestOnceOrRetry_Do(t *testing.T) {
197+
var once OnceOrRetry
198+
var count atomic.Int32
199+
var wg sync.WaitGroup
200+
for i := 0; i < 100; i++ {
201+
wg.Add(1)
202+
go func() {
203+
defer wg.Done()
204+
err := once.Do(func() error {
205+
count.Add(1)
206+
return nil
207+
})
208+
if err != nil {
209+
t.Errorf("OnceOrRetry.Do() error = %v, wantErr %v", err, nil)
210+
}
211+
}()
212+
}
213+
wg.Wait()
214+
215+
if got := count.Load(); got != 1 {
216+
t.Fatal("OnceOrRetry.Do() called more than once")
217+
}
218+
}
219+
220+
func TestOnceOrRetry_Do_Fail(t *testing.T) {
221+
var once OnceOrRetry
222+
var wg sync.WaitGroup
223+
224+
// test failure
225+
for i := 0; i < 100; i++ {
226+
wg.Add(1)
227+
go func(wantErr error) {
228+
defer wg.Done()
229+
err := once.Do(func() error {
230+
return wantErr
231+
})
232+
if err != wantErr {
233+
t.Errorf("OnceOrRetry.Do() error = %v, wantErr %v", err, wantErr)
234+
}
235+
}(errors.New(strconv.Itoa(i)))
236+
}
237+
wg.Wait()
238+
239+
// retry after failure
240+
err := once.Do(func() error {
241+
return nil
242+
})
243+
if err != nil {
244+
t.Fatalf("OnceOrRetry.Do() error = %v, wantErr %v", err, nil)
245+
}
246+
247+
// no retry after success
248+
err = once.Do(func() error {
249+
t.Fatal("OnceOrRetry.Do() called twice")
250+
return nil
251+
})
252+
if err != nil {
253+
t.Fatalf("OnceOrRetry.Do() error = %v, wantErr %v", err, nil)
254+
}
255+
}
256+
257+
func TestOnceOrRetry_Do_Panic(t *testing.T) {
258+
var once OnceOrRetry
259+
260+
// test panic
261+
func() {
262+
defer func() {
263+
if r := recover(); r == nil {
264+
t.Fatal("OnceOrRetry.Do() did not panic")
265+
}
266+
}()
267+
_ = once.Do(func() error {
268+
panic("failed")
269+
})
270+
}()
271+
272+
// retry after panic
273+
err := once.Do(func() error {
274+
return nil
275+
})
276+
if err != nil {
277+
t.Fatalf("OnceOrRetry.Do() error = %v, wantErr %v", err, nil)
278+
}
279+
280+
// no retry after success
281+
err = once.Do(func() error {
282+
t.Fatal("OnceOrRetry.Do() called twice")
283+
return nil
284+
})
285+
if err != nil {
286+
t.Fatalf("OnceOrRetry.Do() error = %v, wantErr %v", err, nil)
287+
}
288+
}

‎registry/remote/credentials/store.go‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -25,8 +25,8 @@ import (
2525
"fmt"
2626
"os"
2727
"path/filepath"
28-
"sync"
2928

29+
"oras.land/oras-go/v2/internal/syncutil"
3030
"oras.land/oras-go/v2/registry/remote/auth"
3131
"oras.land/oras-go/v2/registry/remote/credentials/internal/config"
3232
)
@@ -53,7 +53,7 @@ type DynamicStore struct {
5353
config *config.Config
5454
options StoreOptions
5555
detectedCredsStore string
56-
setCredsStoreOnce sync.Once
56+
setCredsStoreOnce syncutil.OnceOrRetry
5757
}
5858

5959
// StoreOptions provides options for NewStore.
@@ -136,19 +136,19 @@ func (ds *DynamicStore) Get(ctx context.Context, serverAddress string) (auth.Cre
136136
// Put saves credentials into the store for the given server address.
137137
// Put returns ErrPlaintextPutDisabled if native store is not available and
138138
// [StoreOptions].AllowPlaintextPut is set to false.
139-
func (ds *DynamicStore) Put(ctx context.Context, serverAddress string, cred auth.Credential) (returnErr error) {
139+
func (ds *DynamicStore) Put(ctx context.Context, serverAddress string, cred auth.Credential) error {
140140
if err := ds.getStore(serverAddress).Put(ctx, serverAddress, cred); err != nil {
141141
return err
142142
}
143143
// save the detected creds store back to the config file on first put
144-
ds.setCredsStoreOnce.Do(func() {
144+
return ds.setCredsStoreOnce.Do(func() error {
145145
if ds.detectedCredsStore != "" {
146146
if err := ds.config.SetCredentialsStore(ds.detectedCredsStore); err != nil {
147-
returnErr = fmt.Errorf("failed to set credsStore: %w", err)
147+
return fmt.Errorf("failed to set credsStore: %w", err)
148148
}
149149
}
150+
return nil
150151
})
151-
return returnErr
152152
}
153153

154154
// Delete removes credentials from the store for the given server address.

0 commit comments

Comments
 (0)