|
| 1 | +package filehandler |
| 2 | + |
| 3 | +import ( |
| 4 | + "context" |
| 5 | + "fmt" |
| 6 | + "node-agent/pkg/filehandler" |
| 7 | + "sync" |
| 8 | +) |
| 9 | + |
| 10 | +const initFileListLength = 5000 |
| 11 | + |
| 12 | +type InMemoryFileHandler struct { |
| 13 | + mutex sync.RWMutex |
| 14 | + m map[string]*sync.RWMutex |
| 15 | + files map[string]map[string]bool |
| 16 | +} |
| 17 | + |
| 18 | +var _ filehandler.FileHandler = (*InMemoryFileHandler)(nil) |
| 19 | + |
| 20 | +func CreateInMemoryFileHandler() (*InMemoryFileHandler, error) { |
| 21 | + return &InMemoryFileHandler{ |
| 22 | + m: make(map[string]*sync.RWMutex), |
| 23 | + files: make(map[string]map[string]bool, 20), |
| 24 | + }, nil |
| 25 | +} |
| 26 | + |
| 27 | +func (s *InMemoryFileHandler) AddFile(ctx context.Context, bucket, file string) error { |
| 28 | + |
| 29 | + // Acquire a read lock first |
| 30 | + s.mutex.RLock() |
| 31 | + bucketLock, ok := s.m[bucket] |
| 32 | + bucketFiles, okF := s.files[bucket] |
| 33 | + s.mutex.RUnlock() |
| 34 | + |
| 35 | + // If the bucket doesn't exist, acquire a write lock to create the new bucket |
| 36 | + if !ok || !okF { |
| 37 | + s.mutex.Lock() |
| 38 | + // Double-check the bucket's existence to ensure another goroutine didn't already create it |
| 39 | + bucketLock, ok = s.m[bucket] |
| 40 | + if !ok { |
| 41 | + bucketLock = &sync.RWMutex{} |
| 42 | + s.m[bucket] = bucketLock |
| 43 | + } |
| 44 | + |
| 45 | + bucketFiles, okF = s.files[bucket] |
| 46 | + if !okF { |
| 47 | + bucketFiles = make(map[string]bool, initFileListLength) |
| 48 | + s.files[bucket] = bucketFiles |
| 49 | + } |
| 50 | + s.mutex.Unlock() |
| 51 | + } |
| 52 | + |
| 53 | + // Acquire a write lock if the bucket already exists |
| 54 | + bucketLock.Lock() |
| 55 | + defer bucketLock.Unlock() |
| 56 | + |
| 57 | + bucketFiles[file] = true |
| 58 | + |
| 59 | + return nil |
| 60 | +} |
| 61 | + |
| 62 | +func (s *InMemoryFileHandler) Close() { |
| 63 | + // Nothing to do |
| 64 | +} |
| 65 | + |
| 66 | +func shallowCopyMapStringBool(m map[string]bool) map[string]bool { |
| 67 | + if m == nil { |
| 68 | + return nil |
| 69 | + } |
| 70 | + mCopy := make(map[string]bool, len(m)) |
| 71 | + for k, v := range m { |
| 72 | + mCopy[k] = v |
| 73 | + } |
| 74 | + return mCopy |
| 75 | +} |
| 76 | + |
| 77 | +func (s *InMemoryFileHandler) GetFiles(ctx context.Context, bucket string) (map[string]bool, error) { |
| 78 | + s.mutex.RLock() |
| 79 | + bucketLock, ok := s.m[bucket] |
| 80 | + bucketFiles, okFiles := s.files[bucket] |
| 81 | + s.mutex.RUnlock() |
| 82 | + |
| 83 | + if !ok || !okFiles { |
| 84 | + return map[string]bool{}, fmt.Errorf("bucket does not exist for container %s", bucket) |
| 85 | + } |
| 86 | + |
| 87 | + bucketLock.RLock() |
| 88 | + defer bucketLock.RUnlock() |
| 89 | + |
| 90 | + return shallowCopyMapStringBool(bucketFiles), nil |
| 91 | +} |
| 92 | +func (s *InMemoryFileHandler) RemoveBucket(ctx context.Context, bucket string) error { |
| 93 | + |
| 94 | + s.mutex.Lock() |
| 95 | + bucketLock, ok := s.m[bucket] |
| 96 | + if ok { |
| 97 | + bucketLock.Lock() |
| 98 | + defer bucketLock.Unlock() |
| 99 | + } |
| 100 | + |
| 101 | + delete(s.m, bucket) |
| 102 | + delete(s.files, bucket) |
| 103 | + s.mutex.Unlock() |
| 104 | + |
| 105 | + return nil |
| 106 | +} |
0 commit comments