Merge pull request #106 from kubescape/refactor-in-memory-fh

refactor inmemory file handler

Author: amirmalka

pkg/filehandler/v1/inmemory.go

@@ -9,52 +9,49 @@ import (
 
 const initFileListLength = 5000
 
+type filesBucket struct {
+	lock  *sync.RWMutex
+	files map[string]bool
+}
+
 type InMemoryFileHandler struct {
-	mutex sync.RWMutex
-	m     map[string]*sync.RWMutex
-	files map[string]map[string]bool
+	mutex   sync.RWMutex
+	buckets map[string]*filesBucket
 }
 
 var _ filehandler.FileHandler = (*InMemoryFileHandler)(nil)
 
 func CreateInMemoryFileHandler() (*InMemoryFileHandler, error) {
 	return &InMemoryFileHandler{
-		m:     make(map[string]*sync.RWMutex),
-		files: make(map[string]map[string]bool, 20),
+		buckets: make(map[string]*filesBucket, 20),
 	}, nil
 }
 
 func (s *InMemoryFileHandler) AddFile(ctx context.Context, bucket, file string) error {
-
 	// Acquire a read lock first
 	s.mutex.RLock()
-	bucketLock, ok := s.m[bucket]
-	bucketFiles, okF := s.files[bucket]
+	bucketFiles, ok := s.buckets[bucket]
 	s.mutex.RUnlock()
 
 	// If the bucket doesn't exist, acquire a write lock to create the new bucket
-	if !ok || !okF {
+	if !ok {
 		s.mutex.Lock()
 		// Double-check the bucket's existence to ensure another goroutine didn't already create it
-		bucketLock, ok = s.m[bucket]
+		bucketFiles, ok = s.buckets[bucket]
 		if !ok {
-			bucketLock = &sync.RWMutex{}
-			s.m[bucket] = bucketLock
-		}
-
-		bucketFiles, okF = s.files[bucket]
-		if !okF {
-			bucketFiles = make(map[string]bool, initFileListLength)
-			s.files[bucket] = bucketFiles
+			bucketFiles = &filesBucket{
+				lock:  &sync.RWMutex{},
+				files: make(map[string]bool, initFileListLength),
+			}
+			s.buckets[bucket] = bucketFiles
 		}
 		s.mutex.Unlock()
 	}
 
 	// Acquire a write lock if the bucket already exists
-	bucketLock.Lock()
-	defer bucketLock.Unlock()
-
-	bucketFiles[file] = true
+	bucketFiles.lock.Lock()
+	bucketFiles.files[file] = true
+	bucketFiles.lock.Unlock()
 
 	return nil
 }
@@ -76,30 +73,22 @@ func shallowCopyMapStringBool(m map[string]bool) map[string]bool {
 
 func (s *InMemoryFileHandler) GetFiles(ctx context.Context, bucket string) (map[string]bool, error) {
 	s.mutex.RLock()
-	bucketLock, ok := s.m[bucket]
-	bucketFiles, okFiles := s.files[bucket]
+	bucketFiles, ok := s.buckets[bucket]
 	s.mutex.RUnlock()
 
-	if !ok || !okFiles {
+	if !ok {
 		return map[string]bool{}, fmt.Errorf("bucket does not exist for container %s", bucket)
 	}
 
-	bucketLock.RLock()
-	defer bucketLock.RUnlock()
+	bucketFiles.lock.RLock()
+	copy := shallowCopyMapStringBool(bucketFiles.files)
+	bucketFiles.lock.RUnlock()
 
-	return shallowCopyMapStringBool(bucketFiles), nil
+	return copy, nil
 }
 func (s *InMemoryFileHandler) RemoveBucket(ctx context.Context, bucket string) error {
-
 	s.mutex.Lock()
-	bucketLock, ok := s.m[bucket]
-	if ok {
-		bucketLock.Lock()
-		defer bucketLock.Unlock()
-	}
-
-	delete(s.m, bucket)
-	delete(s.files, bucket)
+	delete(s.buckets, bucket)
 	s.mutex.Unlock()
 
 	return nil