merge v0.23.0-rc changes

2024-09-29 19:23:19 +03:00
parent ad92992324
commit 844f18cac3
753 changed files with 85141 additions and 63396 deletions
@@ -1,11 +1,18 @@
 package store

-import "sync"
+import (
+	"encoding/json"
+	"sync"
+)
+
+// @todo remove after https://github.com/golang/go/issues/20135
+const ShrinkThreshold = 200 // the number is arbitrary chosen

 // Store defines a concurrent safe in memory key-value data store.
 type Store[T any] struct {
-	data map[string]T
-	mux  sync.RWMutex
+	data    map[string]T
+	mu      sync.RWMutex
+	deleted int64
 }

 // New creates a new Store[T] instance with a shallow copy of the provided data (if any).
@@ -20,8 +27,8 @@ func New[T any](data map[string]T) *Store[T] {
 // Reset clears the store and replaces the store data with a
 // shallow copy of the provided newData.
 func (s *Store[T]) Reset(newData map[string]T) {
-	s.mux.Lock()
-	defer s.mux.Unlock()
+	s.mu.Lock()
+	defer s.mu.Unlock()

 	if len(newData) > 0 {
 		s.data = make(map[string]T, len(newData))
@@ -31,38 +38,50 @@ func (s *Store[T]) Reset(newData map[string]T) {
 	} else {
 		s.data = make(map[string]T)
 	}
+
+	s.deleted = 0
 }

 // Length returns the current number of elements in the store.
 func (s *Store[T]) Length() int {
-	s.mux.RLock()
-	defer s.mux.RUnlock()
+	s.mu.RLock()
+	defer s.mu.RUnlock()

 	return len(s.data)
 }

 // RemoveAll removes all the existing store entries.
 func (s *Store[T]) RemoveAll() {
-	s.mux.Lock()
-	defer s.mux.Unlock()
-
-	s.data = make(map[string]T)
+	s.Reset(nil)
 }

 // Remove removes a single entry from the store.
 //
 // Remove does nothing if key doesn't exist in the store.
 func (s *Store[T]) Remove(key string) {
-	s.mux.Lock()
-	defer s.mux.Unlock()
+	s.mu.Lock()
+	defer s.mu.Unlock()

 	delete(s.data, key)
+	s.deleted++
+
+	// reassign to a new map so that the old one can be gc-ed because it doesn't shrink
+	//
+	// @todo remove after https://github.com/golang/go/issues/20135
+	if s.deleted >= ShrinkThreshold {
+		newData := make(map[string]T, len(s.data))
+		for k, v := range s.data {
+			newData[k] = v
+		}
+		s.data = newData
+		s.deleted = 0
+	}
 }

 // Has checks if element with the specified key exist or not.
 func (s *Store[T]) Has(key string) bool {
-	s.mux.RLock()
-	defer s.mux.RUnlock()
+	s.mu.RLock()
+	defer s.mu.RUnlock()

 	_, ok := s.data[key]

@@ -73,16 +92,26 @@ func (s *Store[T]) Has(key string) bool {
 //
 // If key is not set, the zero T value is returned.
 func (s *Store[T]) Get(key string) T {
-	s.mux.RLock()
-	defer s.mux.RUnlock()
+	s.mu.RLock()
+	defer s.mu.RUnlock()

 	return s.data[key]
 }

+// GetOk is similar to Get but returns also a boolean indicating whether the key exists or not.
+func (s *Store[T]) GetOk(key string) (T, bool) {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+
+	v, ok := s.data[key]
+
+	return v, ok
+}
+
 // GetAll returns a shallow copy of the current store data.
 func (s *Store[T]) GetAll() map[string]T {
-	s.mux.RLock()
-	defer s.mux.RUnlock()
+	s.mu.RLock()
+	defer s.mu.RUnlock()

 	var clone = make(map[string]T, len(s.data))

@@ -93,10 +122,24 @@ func (s *Store[T]) GetAll() map[string]T {
 	return clone
 }

+// Values returns a slice with all of the current store values.
+func (s *Store[T]) Values() []T {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+
+	var values = make([]T, 0, len(s.data))
+
+	for _, v := range s.data {
+		values = append(values, v)
+	}
+
+	return values
+}
+
 // Set sets (or overwrite if already exist) a new value for key.
 func (s *Store[T]) Set(key string, value T) {
-	s.mux.Lock()
-	defer s.mux.Unlock()
+	s.mu.Lock()
+	defer s.mu.Unlock()

 	if s.data == nil {
 		s.data = make(map[string]T)
@@ -105,16 +148,34 @@ func (s *Store[T]) Set(key string, value T) {
 	s.data[key] = value
 }

+// GetOrSet retrieves a single existing value for the provided key
+// or stores a new one if it doesn't exist.
+func (s *Store[T]) GetOrSet(key string, setFunc func() T) T {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	if s.data == nil {
+		s.data = make(map[string]T)
+	}
+
+	v, ok := s.data[key]
+	if !ok {
+		v = setFunc()
+		s.data[key] = v
+	}
+
+	return v
+}
+
 // SetIfLessThanLimit sets (or overwrite if already exist) a new value for key.
 //
 // This method is similar to Set() but **it will skip adding new elements**
 // to the store if the store length has reached the specified limit.
 // false is returned if maxAllowedElements limit is reached.
 func (s *Store[T]) SetIfLessThanLimit(key string, value T, maxAllowedElements int) bool {
-	s.mux.Lock()
-	defer s.mux.Unlock()
+	s.mu.Lock()
+	defer s.mu.Unlock()

-	// init map if not already
 	if s.data == nil {
 		s.data = make(map[string]T)
 	}
@@ -132,3 +193,33 @@ func (s *Store[T]) SetIfLessThanLimit(key string, value T, maxAllowedElements in

 	return true
 }
+
+// UnmarshalJSON implements [json.Unmarshaler] and imports the
+// provided JSON data into the store.
+//
+// The store entries that match with the ones from the data will be overwritten with the new value.
+func (s *Store[T]) UnmarshalJSON(data []byte) error {
+	raw := map[string]T{}
+	if err := json.Unmarshal(data, &raw); err != nil {
+		return err
+	}
+
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	if s.data == nil {
+		s.data = make(map[string]T)
+	}
+
+	for k, v := range raw {
+		s.data[k] = v
+	}
+
+	return nil
+}
+
+// MarshalJSON implements [json.Marshaler] and export the current
+// store data into valid JSON.
+func (s *Store[T]) MarshalJSON() ([]byte, error) {
+	return json.Marshal(s.GetAll())
+}
@@ -3,6 +3,8 @@ package store_test
 import (
 	"bytes"
 	"encoding/json"
+	"slices"
+	"strconv"
 	"testing"

 	"github.com/pocketbase/pocketbase/tools/store"
@@ -137,11 +139,41 @@ func TestGet(t *testing.T) {
 		{"missing", 0}, // should auto fallback to the zero value
 	}

-	for i, scenario := range scenarios {
-		val := s.Get(scenario.key)
-		if val != scenario.expect {
-			t.Errorf("(%d) Expected %v, got %v", i, scenario.expect, val)
-		}
+	for _, scenario := range scenarios {
+		t.Run(scenario.key, func(t *testing.T) {
+			val := s.Get(scenario.key)
+			if val != scenario.expect {
+				t.Fatalf("Expected %v, got %v", scenario.expect, val)
+			}
+		})
+	}
+}
+
+func TestGetOk(t *testing.T) {
+	s := store.New(map[string]int{"test1": 0, "test2": 1})
+
+	scenarios := []struct {
+		key         string
+		expectValue int
+		expectOk    bool
+	}{
+		{"test1", 0, true},
+		{"test2", 1, true},
+		{"missing", 0, false}, // should auto fallback to the zero value
+	}
+
+	for _, scenario := range scenarios {
+		t.Run(scenario.key, func(t *testing.T) {
+			val, ok := s.GetOk(scenario.key)
+
+			if ok != scenario.expectOk {
+				t.Fatalf("Expected ok %v, got %v", scenario.expectOk, ok)
+			}
+
+			if val != scenario.expectValue {
+				t.Fatalf("Expected %v, got %v", scenario.expectValue, val)
+			}
+		})
 	}
 }

@@ -173,6 +205,27 @@ func TestGetAll(t *testing.T) {
 	}
 }

+func TestValues(t *testing.T) {
+	data := map[string]int{
+		"a": 1,
+		"b": 2,
+	}
+
+	values := store.New(data).Values()
+
+	expected := []int{1, 2}
+
+	if len(values) != len(expected) {
+		t.Fatalf("Expected %d values, got %d", len(expected), len(values))
+	}
+
+	for _, v := range expected {
+		if !slices.Contains(values, v) {
+			t.Fatalf("Missing value %v in\n%v", v, values)
+		}
+	}
+}
+
 func TestSet(t *testing.T) {
 	s := store.Store[int]{}

@@ -196,6 +249,37 @@ func TestSet(t *testing.T) {
 	}
 }

+func TestGetOrSet(t *testing.T) {
+	s := store.New(map[string]int{
+		"test1": 0,
+		"test2": 1,
+		"test3": 3,
+	})
+
+	scenarios := []struct {
+		key      string
+		value    int
+		expected int
+	}{
+		{"test2", 20, 1},
+		{"test3", 2, 3},
+		{"test_new", 20, 20},
+		{"test_new", 50, 20}, // should return the previously inserted value
+	}
+
+	for _, scenario := range scenarios {
+		t.Run(scenario.key, func(t *testing.T) {
+			result := s.GetOrSet(scenario.key, func() int {
+				return scenario.value
+			})
+
+			if result != scenario.expected {
+				t.Fatalf("Expected %v, got %v", scenario.expected, result)
+			}
+		})
+	}
+}
+
 func TestSetIfLessThanLimit(t *testing.T) {
 	s := store.Store[int]{}

@@ -230,3 +314,77 @@ func TestSetIfLessThanLimit(t *testing.T) {
 		}
 	}
 }
+
+func TestUnmarshalJSON(t *testing.T) {
+	s := store.Store[string]{}
+	s.Set("b", "old")   // should be overwritten
+	s.Set("c", "test3") // ensures that the old values are not removed
+
+	raw := []byte(`{"a":"test1", "b":"test2"}`)
+	if err := json.Unmarshal(raw, &s); err != nil {
+		t.Fatal(err)
+	}
+
+	if v := s.Get("a"); v != "test1" {
+		t.Fatalf("Expected store.a to be %q, got %q", "test1", v)
+	}
+
+	if v := s.Get("b"); v != "test2" {
+		t.Fatalf("Expected store.b to be %q, got %q", "test2", v)
+	}
+
+	if v := s.Get("c"); v != "test3" {
+		t.Fatalf("Expected store.c to be %q, got %q", "test3", v)
+	}
+}
+
+func TestMarshalJSON(t *testing.T) {
+	s := &store.Store[string]{}
+	s.Set("a", "test1")
+	s.Set("b", "test2")
+
+	expected := []byte(`{"a":"test1", "b":"test2"}`)
+
+	result, err := json.Marshal(s)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if bytes.Equal(result, expected) {
+		t.Fatalf("Expected\n%s\ngot\n%s", expected, result)
+	}
+}
+
+func TestShrink(t *testing.T) {
+	s := &store.Store[int]{}
+
+	total := 1000
+
+	for i := 0; i < total; i++ {
+		s.Set(strconv.Itoa(i), i)
+	}
+
+	if s.Length() != total {
+		t.Fatalf("Expected %d items, got %d", total, s.Length())
+	}
+
+	// trigger map "shrink"
+	for i := 0; i < store.ShrinkThreshold; i++ {
+		s.Remove(strconv.Itoa(i))
+	}
+
+	// ensure that after the deletion, the new map was copied properly
+	if s.Length() != total-store.ShrinkThreshold {
+		t.Fatalf("Expected %d items, got %d", total-store.ShrinkThreshold, s.Length())
+	}
+
+	for k := range s.GetAll() {
+		kInt, err := strconv.Atoi(k)
+		if err != nil {
+			t.Fatalf("failed to convert %s into int: %v", k, err)
+		}
+		if kInt < store.ShrinkThreshold {
+			t.Fatalf("Key %q should have been deleted", k)
+		}
+	}
+}