cache.go 2.8 KB

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