Repository: vllm-project/aibrix
File: pkg/plugins/gateway/prefixcacheindexer/hash.go (Go)

Code before:
	modelToPods    map[string]map[string]time.Time // model_name: map[pod_name]pod_last_access_time
	lastAccessTime time.Time                       // block_last_access_time
}

func NewPrefixHashTable() PrefixCacheIndexer {
	r := rand.New(rand.NewSource(time.Now().Unix()))
	seed := r.Uint64()
	instance := &PrefixHashTable{
		blocks: map[uint64]Block{},
		hash:   xxhash.NewWithSeed(seed),
		seed:   seed,
	}

	ticker := time.NewTicker(prefixCacheEvictionInterval)
	go func() {
		for range ticker.C {
			instance.Evict(time.Now())
		}
	}()

	return instance
}

// returns matchedTokens, unMatchedTokens, matchedPods
func (c *PrefixHashTable) MatchPrefix(tokens []byte, model string, pods []*v1.Pod) ([]byte, []byte, []*v1.Pod) {
	c.mu.Lock()
	defer c.mu.Unlock()
	var block, lastMatchedBlock Block
	var ok bool
	var lastTokenMatchIndex int

	for i := 0; i < len(tokens); i += prefixCacheBlockSize {
		end := i + prefixCacheBlockSize
		if end > len(tokens) {
			end = len(tokens)
		}

		_, _ = c.hash.Write(tokens[i:end])
		prefixHash := c.hash.Sum64()
		c.hash.ResetWithSeed(c.seed)

		block, ok = c.blocks[prefixHash]
		if !ok || len(block.modelToPods[model]) == 0 || len(matchPods(block.modelToPods[model], pods)) == 0 {
			lastTokenMatchIndex = i
			break
		}

		lastTokenMatchIndex = end
		lastMatchedBlock = block
		block.lastAccessTime = time.Now()
		c.blocks[prefixHash] = block

Review comment (bug):
@vie-serendipity I assumed LRU store is thread-safe and removing this lock was fine. While working on PR #933 I came across the bug for concurrent read/write for blocks.

I have added the lock back, if you get a chance, can you make LRU store thread-safe? Thank you.

Code after:
			klog.Infof("invalid AIBRIX_PREFIX_CACHE_EVICTION_DURATION_MINS: %s, falling back to default", value)
		} else {
			klog.Infof("using AIBRIX_PREFIX_CACHE_EVICTION_DURATION_MINS env value for prefix cache eviction duration: %d ms", intValue)
			return time.Duration(intValue) * time.Minute
		}
	}
	klog.Infof("using default prefix cache eviction duration: %d mins", defaultPrefixCacheEvictionDurationInMins)
	return defaultPrefixCacheEvictionDurationInMins * time.Minute
}

type PrefixHashTable struct {
	mu    sync.RWMutex
	hash  *xxhash.Digest
	seed  uint64
	store cache.Store[uint64, Block]
}

type Block struct {
	modelToPods map[string]map[string]time.Time // model_name: map[pod_name]pod_last_access_time
}

func NewPrefixHashTable() PrefixCacheIndexer {
	r := rand.New(rand.NewSource(time.Now().Unix()))
	seed := r.Uint64()
	instance := &PrefixHashTable{
		hash: xxhash.NewWithSeed(seed),
		seed: seed,
		store: cache.NewLRUStore[uint64, Block](prefixCacheBlockNumber,
			prefixCacheEvictionDuration,
			prefixCacheEvictionInterval,
			func() time.Time { return time.Now() }),
	}

	return instance
}

// returns matchedTokens, unMatchedTokens, matchedPods
func (c *PrefixHashTable) MatchPrefix(tokens []byte, model string, pods []*v1.Pod) ([]byte, []byte, []*v1.Pod) {
	var block, lastMatchedBlock Block
	var ok bool
	var lastTokenMatchIndex int

	for i := 0; i < len(tokens); i += prefixCacheBlockSize {
		end := i + prefixCacheBlockSize
		if end > len(tokens) {
			end = len(tokens)
		}

		c.mu.Lock()
		_, _ = c.hash.Write(tokens[i:end])
		prefixHash := c.hash.Sum64()