github.com/crowdsecurity/crowdsec@v1.6.1/pkg/cache/cache.go (about)

     1  package cache
     2  
     3  import (
     4  	"errors"
     5  	"time"
     6  
     7  	"github.com/bluele/gcache"
     8  	"github.com/prometheus/client_golang/prometheus"
     9  	log "github.com/sirupsen/logrus"
    10  
    11  	"github.com/crowdsecurity/crowdsec/pkg/types"
    12  )
    13  
    14  var Caches []gcache.Cache
    15  var CacheNames []string
    16  var CacheConfig []CacheCfg
    17  
    18  /*prometheus*/
    19  var CacheMetrics = prometheus.NewGaugeVec(
    20  	prometheus.GaugeOpts{
    21  		Name: "cs_cache_size",
    22  		Help: "Entries per cache.",
    23  	},
    24  	[]string{"name", "type"},
    25  )
    26  
    27  // UpdateCacheMetrics is called directly by the prom handler
    28  func UpdateCacheMetrics() {
    29  	CacheMetrics.Reset()
    30  	for i, name := range CacheNames {
    31  		CacheMetrics.With(prometheus.Labels{"name": name, "type": CacheConfig[i].Strategy}).Set(float64(Caches[i].Len(false)))
    32  	}
    33  }
    34  
    35  type CacheCfg struct {
    36  	Name     string
    37  	Size     int
    38  	TTL      time.Duration
    39  	Strategy string
    40  	LogLevel *log.Level
    41  	Logger   *log.Entry
    42  }
    43  
    44  func CacheInit(cfg CacheCfg) error {
    45  
    46  	for _, name := range CacheNames {
    47  		if name == cfg.Name {
    48  			log.Infof("Cache %s already exists", cfg.Name)
    49  		}
    50  	}
    51  	//get a default logger
    52  	if cfg.LogLevel == nil {
    53  		cfg.LogLevel = new(log.Level)
    54  		*cfg.LogLevel = log.InfoLevel
    55  	}
    56  	var clog = log.New()
    57  	if err := types.ConfigureLogger(clog); err != nil {
    58  		log.Fatalf("While creating cache logger : %s", err)
    59  	}
    60  	clog.SetLevel(*cfg.LogLevel)
    61  	cfg.Logger = clog.WithFields(log.Fields{
    62  		"cache": cfg.Name,
    63  	})
    64  
    65  	tmpCache := gcache.New(cfg.Size)
    66  	switch cfg.Strategy {
    67  	case "LRU":
    68  		tmpCache = tmpCache.LRU()
    69  	case "LFU":
    70  		tmpCache = tmpCache.LFU()
    71  	case "ARC":
    72  		tmpCache = tmpCache.ARC()
    73  	default:
    74  		cfg.Strategy = "LRU"
    75  		tmpCache = tmpCache.LRU()
    76  
    77  	}
    78  
    79  	CTICache := tmpCache.Build()
    80  	Caches = append(Caches, CTICache)
    81  	CacheNames = append(CacheNames, cfg.Name)
    82  	CacheConfig = append(CacheConfig, cfg)
    83  
    84  	return nil
    85  }
    86  
    87  func SetKey(cacheName string, key string, value string, expiration *time.Duration) error {
    88  
    89  	for i, name := range CacheNames {
    90  		if name == cacheName {
    91  			if expiration == nil {
    92  				expiration = &CacheConfig[i].TTL
    93  			}
    94  			CacheConfig[i].Logger.Debugf("Setting key %s to %s with expiration %v", key, value, *expiration)
    95  			if err := Caches[i].SetWithExpire(key, value, *expiration); err != nil {
    96  				CacheConfig[i].Logger.Warningf("While setting key %s in cache %s: %s", key, cacheName, err)
    97  			}
    98  		}
    99  	}
   100  	return nil
   101  }
   102  
   103  func GetKey(cacheName string, key string) (string, error) {
   104  	for i, name := range CacheNames {
   105  		if name == cacheName {
   106  			if value, err := Caches[i].Get(key); err != nil {
   107  				//do not warn or log if key not found
   108  				if errors.Is(err, gcache.KeyNotFoundError) {
   109  					return "", nil
   110  				}
   111  				CacheConfig[i].Logger.Warningf("While getting key %s in cache %s: %s", key, cacheName, err)
   112  				return "", err
   113  			} else {
   114  				return value.(string), nil
   115  			}
   116  		}
   117  	}
   118  	log.Warningf("Cache %s not found", cacheName)
   119  	return "", nil
   120  }