[BUG] Fix non thread safe token cache behavior

[pmahindrakar-oss]

Describe the bug

Currently the client auth interceptor used by multiple flyte services faces an issue when multiple parallel attempts are made to get token as the underlying token source implementation is not thread safe

https://github.com/flyteorg/flyte/blob/master/flyteidl/clients/go/admin/auth_interceptor.go#L145

Eg : https://github.com/flyteorg/flyte/blob/master/flyteidl/clients/go/admin/cache/token_cache_inmemory.go

Expected behavior

The implementation of token cache should be thread safe

Additional context to reproduce

No response

Screenshots

No response

Are you sure this issue hasn't been raised already?

• Yes

Have you read the Code of Conduct?

• Yes

[runllm]

The current implementation of the token cache in Flyte is not thread-safe, which can cause issues when multiple parallel attempts are made to get a token. This is evident in the TokenCacheInMemoryProvider class, which does not use any synchronization mechanisms to ensure thread safety.

To make the token cache thread-safe, you can use synchronization primitives like mutexes. Here's an example of how you can modify the TokenCacheInMemoryProvider to be thread-safe:

package cache

import (
	"fmt"
	"sync"

	"golang.org/x/oauth2"
)

type TokenCacheInMemoryProvider struct {
	token *oauth2.Token
	mu    sync.Mutex
}

func (t *TokenCacheInMemoryProvider) SaveToken(token *oauth2.Token) error {
	t.mu.Lock()
	defer t.mu.Unlock()
	t.token = token
	return nil
}

func (t *TokenCacheInMemoryProvider) GetToken() (*oauth2.Token, error) {
	t.mu.Lock()
	defer t.mu.Unlock()
	if t.token == nil {
		return nil, fmt.Errorf("cannot find token in cache")
	}
	return t.token, nil
}

Would you like more details on this implementation?

Sources:

• Does Flyte support the caching of M2M authorization tokens?
• Hi there we ve been having some challenges with the token ca
• Core feature pyflyte should store tokens after DeviceFlow auth

---

This is an AI-generated response and your feedback is appreciated! Please leave a 👍 if this is helpful and 👎 if it is not.

Mention @RunLLM to ask follow-up questions. See runllm.com for more!