116 lines
2.7 KiB
Go
116 lines
2.7 KiB
Go
package syncer
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"gemini-balancer/internal/store"
|
|
"log"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
// LoaderFunc
|
|
type LoaderFunc[T any] func() (T, error)
|
|
|
|
// CacheSyncer
|
|
type CacheSyncer[T any] struct {
|
|
mu sync.RWMutex
|
|
cache T
|
|
loader LoaderFunc[T]
|
|
store store.Store
|
|
channelName string
|
|
stopChan chan struct{}
|
|
wg sync.WaitGroup
|
|
}
|
|
|
|
// NewCacheSyncer
|
|
func NewCacheSyncer[T any](
|
|
loader LoaderFunc[T],
|
|
store store.Store,
|
|
channelName string,
|
|
) (*CacheSyncer[T], error) {
|
|
s := &CacheSyncer[T]{
|
|
loader: loader,
|
|
store: store,
|
|
channelName: channelName,
|
|
stopChan: make(chan struct{}),
|
|
}
|
|
if err := s.reload(); err != nil {
|
|
return nil, fmt.Errorf("initial load for %s failed: %w", channelName, err)
|
|
}
|
|
s.wg.Add(1)
|
|
go s.listenForUpdates()
|
|
return s, nil
|
|
}
|
|
|
|
// Get, Invalidate, Stop, reload 方法 .
|
|
func (s *CacheSyncer[T]) Get() T {
|
|
s.mu.RLock()
|
|
defer s.mu.RUnlock()
|
|
return s.cache
|
|
}
|
|
|
|
func (s *CacheSyncer[T]) Invalidate() error {
|
|
log.Printf("INFO: Publishing invalidation notification on channel '%s'", s.channelName)
|
|
return s.store.Publish(context.Background(), s.channelName, []byte("reload"))
|
|
}
|
|
|
|
func (s *CacheSyncer[T]) Stop() {
|
|
close(s.stopChan)
|
|
s.wg.Wait()
|
|
log.Printf("INFO: CacheSyncer for channel '%s' stopped.", s.channelName)
|
|
}
|
|
|
|
func (s *CacheSyncer[T]) reload() error {
|
|
log.Printf("INFO: Reloading cache for channel '%s'...", s.channelName)
|
|
newData, err := s.loader()
|
|
if err != nil {
|
|
log.Printf("ERROR: Failed to reload cache for '%s': %v", s.channelName, err)
|
|
return err
|
|
}
|
|
s.mu.Lock()
|
|
s.cache = newData
|
|
s.mu.Unlock()
|
|
log.Printf("INFO: Cache for channel '%s' reloaded successfully.", s.channelName)
|
|
return nil
|
|
}
|
|
|
|
// listenForUpdates ...
|
|
func (s *CacheSyncer[T]) listenForUpdates() {
|
|
defer s.wg.Done()
|
|
for {
|
|
select {
|
|
case <-s.stopChan:
|
|
return
|
|
default:
|
|
}
|
|
|
|
subscription, err := s.store.Subscribe(context.Background(), s.channelName)
|
|
if err != nil {
|
|
log.Printf("ERROR: Failed to subscribe to '%s', retrying in 5s: %v", s.channelName, err)
|
|
time.Sleep(5 * time.Second)
|
|
continue
|
|
}
|
|
log.Printf("INFO: Subscribed to channel '%s' for cache invalidation.", s.channelName)
|
|
|
|
subscriberLoop:
|
|
for {
|
|
select {
|
|
case _, ok := <-subscription.Channel():
|
|
if !ok {
|
|
log.Printf("WARN: Subscription channel '%s' closed, will re-subscribe.", s.channelName)
|
|
break subscriberLoop
|
|
}
|
|
log.Printf("INFO: Received invalidation notification on '%s', reloading cache.", s.channelName)
|
|
if err := s.reload(); err != nil {
|
|
log.Printf("ERROR: Failed to reload cache for '%s' after notification: %v", s.channelName, err)
|
|
}
|
|
case <-s.stopChan:
|
|
subscription.Close()
|
|
return
|
|
}
|
|
}
|
|
subscription.Close()
|
|
}
|
|
}
|