Use sync.Map

handler/transport/oauth2_req_auth.go

@@ -13,15 +13,12 @@ import (
 
 var _ http.RoundTripper = &OAuth2ReqAuth{}
 
-type locks map[string]*sync.Mutex
-
 // OAuth2ReqAuth represents the transport <OAuth2ReqAuth> object.
 type OAuth2ReqAuth struct {
 	oauth2   *OAuth2
 	config   *config.OAuth2ReqAuth
 	memStore *cache.MemoryStore
-	locks    locks
-	locksMu  *sync.Mutex
+	locks    sync.Map
 	next     http.RoundTripper
 }
 
@@ -37,8 +34,7 @@ func NewOAuth2ReqAuth(conf *config.OAuth2ReqAuth, memStore *cache.MemoryStore,
 		config:   conf,
 		oauth2:   oauth2,
 		memStore: memStore,
-		locks:    make(locks),
-		locksMu:  &sync.Mutex{},
+		locks:    sync.Map{},
 		next:     next,
 	}, nil
 }
@@ -58,38 +54,33 @@ func (oa *OAuth2ReqAuth) RoundTrip(req *http.Request) (*http.Response, error) {
 		return oa.next.RoundTrip(req)
 	}
 
-	oa.locksMu.Lock()
-	if _, ok := oa.locks[requestConfig.StorageKey]; !ok {
-		mu := sync.Mutex{}
-		oa.locks[requestConfig.StorageKey] = &mu
-	}
-	oa.locksMu.Unlock()
-
-	oa.locks[requestConfig.StorageKey].Lock()
+	value, _ := oa.locks.LoadOrStore(requestConfig.StorageKey, &sync.Mutex{})
+	mutex := value.(*sync.Mutex)
+	mutex.Lock()
 
 	if token, terr := oa.readAccessToken(requestConfig.StorageKey); terr != nil {
-		oa.locks[requestConfig.StorageKey].Unlock()
+		mutex.Unlock()
 		return nil, errors.Backend.Label(oa.config.BackendName).Message("token read error").With(terr)
 	} else if token != "" {
-		oa.locks[requestConfig.StorageKey].Unlock()
+		mutex.Unlock()
 		req.Header.Set("Authorization", "Bearer "+token)
 
 		return oa.next.RoundTrip(req)
 	}
 
 	tokenResponse, err := oa.oauth2.RequestToken(req.Context(), requestConfig)
 	if err != nil {
-		oa.locks[requestConfig.StorageKey].Unlock()
+		mutex.Unlock()
 		return nil, errors.Backend.Label(oa.config.BackendName).With(err)
 	}
 
 	token, err := oa.updateAccessToken(tokenResponse, requestConfig.StorageKey)
 	if err != nil {
-		oa.locks[requestConfig.StorageKey].Unlock()
+		mutex.Unlock()
 		return nil, errors.Backend.Label(oa.config.BackendName).Message("token update error").With(err)
 	}
 
-	oa.locks[requestConfig.StorageKey].Unlock()
+	mutex.Unlock()
 
 	req.Header.Set("Authorization", "Bearer "+token)