github.com/castai/kvisor@v1.7.1-0.20240516114728-b3572a2607b5/pkg/blobscache/server.go (about)

     1  package blobscache
     2  
     3  import (
     4  	"net/http"
     5  
     6  	"github.com/castai/kvisor/pkg/logging"
     7  	lru "github.com/hashicorp/golang-lru"
     8  	"github.com/labstack/echo/v4"
     9  )
    10  
    11  func NewServer(log *logging.Logger) *Server {
    12  	return &Server{
    13  		log:        log.WithField("component", "blobscache"),
    14  		blobsCache: newMemoryBlobsCacheStore(log),
    15  	}
    16  }
    17  
    18  type Server struct {
    19  	log        *logging.Logger
    20  	blobsCache blobsCacheStore
    21  }
    22  
    23  func (s *Server) RegisterHandlers(e *echo.Echo) {
    24  	e.POST("/blobscache/PutBlob", s.putBlob)
    25  	e.POST("/blobscache/GetBlob", s.getBlob)
    26  }
    27  
    28  func (s *Server) putBlob(c echo.Context) error {
    29  	var req PubBlobRequest
    30  	if err := c.Bind(&req); err != nil {
    31  		return err
    32  	}
    33  	s.blobsCache.putBlob(req.Key, req.Blob)
    34  	return c.NoContent(http.StatusOK)
    35  }
    36  
    37  func (s *Server) getBlob(c echo.Context) error {
    38  	var req GetBlobRequest
    39  	if err := c.Bind(&req); err != nil {
    40  		return err
    41  	}
    42  
    43  	blob, found := s.blobsCache.getBlob(req.Key)
    44  	if !found {
    45  		return c.NoContent(http.StatusNotFound)
    46  	}
    47  
    48  	resp := GetBlobResponse{Blob: blob}
    49  	return c.JSON(http.StatusOK, resp)
    50  }
    51  
    52  type blobsCacheStore interface {
    53  	putBlob(key string, blob []byte)
    54  	getBlob(key string) ([]byte, bool)
    55  }
    56  
    57  func newMemoryBlobsCacheStore(log *logging.Logger) blobsCacheStore {
    58  	// One large blob json size is around 16KB, so we should use max 32MB of extra memory.
    59  	cache, _ := lru.New(2000)
    60  	return &memoryBlobsCacheStore{
    61  		log:   log,
    62  		cache: cache,
    63  	}
    64  }
    65  
    66  type memoryBlobsCacheStore struct {
    67  	log   *logging.Logger
    68  	cache *lru.Cache
    69  }
    70  
    71  func (c *memoryBlobsCacheStore) putBlob(key string, blob []byte) {
    72  	c.log.Debugf("adding image blob to cache, current cache size=%d", c.cache.Len())
    73  	evicted := c.cache.Add(key, blob)
    74  	if evicted {
    75  		c.log.Info("evicted old image blob cache entry")
    76  	}
    77  }
    78  
    79  func (c *memoryBlobsCacheStore) getBlob(key string) ([]byte, bool) {
    80  	val, ok := c.cache.Get(key)
    81  	if !ok {
    82  		return nil, false
    83  	}
    84  	return val.([]byte), true
    85  }