refactor: simplify Once code

Signed-off-by: Shiwei Zhang <shizh@microsoft.com>

internal/syncutil/once.go

@@ -15,10 +15,14 @@ limitations under the License.
 
 package syncutil
 
-import "context"
+import (
+	"context"
+	"sync"
+	"sync/atomic"
+)
 
 // Once is an object that will perform exactly one action.
-// Unlike sync.Once, this Once allowes the action to have return values.
+// Unlike sync.Once, this Once allows the action to have return values.
 type Once struct {
 	result interface{}
 	err    error
@@ -68,3 +72,31 @@ func (o *Once) Do(ctx context.Context, f func() (interface{}, error)) (bool, int
 		}
 	}
 }
+
+// OnceOrRetry is an object that will perform exactly one success action.
+type OnceOrRetry struct {
+	done atomic.Bool
+	lock sync.Mutex
+}
+
+// OnceOrRetry calls the function f if and only if Do is being called for the
+// first time for this instance of Once or all previous calls to Do are failed.
+func (o *OnceOrRetry) Do(f func() error) error {
+	// fast path
+	if o.done.Load() {
+		return nil
+	}
+
+	// slow path
+	o.lock.Lock()
+	defer o.lock.Unlock()
+
+	if o.done.Load() {
+		return nil
+	}
+	if err := f(); err != nil {
+		return err
+	}
+	o.done.Store(true)
+	return nil
+}

registry/remote/credentials/store.go

@@ -53,7 +53,7 @@ type DynamicStore struct {
 	config             *config.Config
 	options            StoreOptions
 	detectedCredsStore string
-	setCredsStoreOnce  func() error
+	setCredsStoreOnce  syncutil.OnceOrRetry
 }
 
 // StoreOptions provides options for NewStore.
@@ -107,14 +107,6 @@ func NewStore(configPath string, opts StoreOptions) (*DynamicStore, error) {
 		// no authentication configured, detect the default credentials store
 		ds.detectedCredsStore = getDefaultHelperSuffix()
 	}
-	ds.setCredsStoreOnce = syncutil.OnceRetryOnError(func() error {
-		if ds.detectedCredsStore != "" {
-			if err := ds.config.SetCredentialsStore(ds.detectedCredsStore); err != nil {
-				return fmt.Errorf("failed to set credsStore: %w", err)
-			}
-		}
-		return nil
-	})
 	return ds, nil
 }
 
@@ -149,7 +141,14 @@ func (ds *DynamicStore) Put(ctx context.Context, serverAddress string, cred auth
 		return err
 	}
 	// save the detected creds store back to the config file on first put
-	return ds.setCredsStoreOnce()
+	return ds.setCredsStoreOnce.Do(func() error {
+		if ds.detectedCredsStore != "" {
+			if err := ds.config.SetCredentialsStore(ds.detectedCredsStore); err != nil {
+				return fmt.Errorf("failed to set credsStore: %w", err)
+			}
+		}
+		return nil
+	})
 }
 
 // Delete removes credentials from the store for the given server address.