github.com/pyroscope-io/pyroscope@v0.37.3-0.20230725203016-5f6947968bd0/pkg/service/auth_cache.go (about)

     1  package service
     2  
     3  import (
     4  	"context"
     5  	"time"
     6  
     7  	lru "github.com/hashicorp/golang-lru"
     8  
     9  	"github.com/pyroscope-io/pyroscope/pkg/model"
    10  )
    11  
    12  type CachingAuthService struct {
    13  	AuthService
    14  	cache *cache
    15  }
    16  
    17  type CachingAuthServiceConfig struct {
    18  	Size int
    19  	TTL  time.Duration
    20  }
    21  
    22  func NewCachingAuthService(authService AuthService, c CachingAuthServiceConfig) CachingAuthService {
    23  	cas := CachingAuthService{AuthService: authService}
    24  	if c.Size > 0 && c.TTL > 0 {
    25  		cas.cache = newCache(c.Size, c.TTL)
    26  	}
    27  	return cas
    28  }
    29  
    30  func (svc CachingAuthService) APIKeyFromToken(ctx context.Context, t string) (model.APIKey, error) {
    31  	if svc.cache != nil {
    32  		return svc.cachedAPIKeyFromToken(ctx, t)
    33  	}
    34  	return svc.AuthService.APIKeyFromToken(ctx, t)
    35  }
    36  
    37  func (svc CachingAuthService) cachedAPIKeyFromToken(ctx context.Context, t string) (model.APIKey, error) {
    38  	if v, ok := svc.cache.get(t); ok {
    39  		switch x := v.(type) {
    40  		case error:
    41  			return model.APIKey{}, x
    42  		case model.APIKey:
    43  			return x, nil
    44  		}
    45  	}
    46  	k, err := svc.AuthService.APIKeyFromToken(ctx, t)
    47  	if err != nil {
    48  		svc.cache.put(t, err)
    49  		return k, err
    50  	}
    51  	svc.cache.put(t, k)
    52  	return k, err
    53  }
    54  
    55  func (svc CachingAuthService) PutAPIKey(t string, k model.APIKey) {
    56  	if svc.cache != nil {
    57  		svc.cache.put(t, k)
    58  	}
    59  }
    60  
    61  func (svc CachingAuthService) DeleteAPIKey(t string) {
    62  	if svc.cache != nil {
    63  		svc.cache.c.Remove(t)
    64  	}
    65  }
    66  
    67  // TODO(kolesnikovae): Move to a separate package.
    68  
    69  type cache struct {
    70  	ttl time.Duration
    71  	c   *lru.Cache
    72  }
    73  
    74  type cachedItem struct {
    75  	value   interface{}
    76  	created time.Time
    77  }
    78  
    79  func newCache(size int, ttl time.Duration) *cache {
    80  	c := cache{ttl: ttl}
    81  	c.c, _ = lru.New(size)
    82  	return &c
    83  }
    84  
    85  func (c *cache) put(k string, v interface{}) {
    86  	c.c.Add(k, cachedItem{
    87  		created: time.Now(),
    88  		value:   v,
    89  	})
    90  }
    91  
    92  func (c *cache) get(k string) (interface{}, bool) {
    93  	x, found := c.c.Get(k)
    94  	if !found {
    95  		return nil, false
    96  	}
    97  	if v, ok := x.(cachedItem); ok && time.Since(v.created) < c.ttl {
    98  		return v.value, ok
    99  	}
   100  	c.c.Remove(k) // Expired.
   101  	return nil, false
   102  }