merge v0.23.0-rc changes

tools/subscriptions/broker.go

@@ -2,47 +2,41 @@ package subscriptions
 
 import (
 	"fmt"
-	"sync"
+
+	"github.com/pocketbase/pocketbase/tools/list"
+	"github.com/pocketbase/pocketbase/tools/store"
 )
 
 // Broker defines a struct for managing subscriptions clients.
 type Broker struct {
-	clients map[string]Client
-	mux     sync.RWMutex
+	store *store.Store[Client]
 }
 
 // NewBroker initializes and returns a new Broker instance.
 func NewBroker() *Broker {
 	return &Broker{
-		clients: make(map[string]Client),
+		store: store.New[Client](nil),
 	}
 }
 
 // Clients returns a shallow copy of all registered clients indexed
 // with their connection id.
 func (b *Broker) Clients() map[string]Client {
-	b.mux.RLock()
-	defer b.mux.RUnlock()
+	return b.store.GetAll()
+}
 
-	copy := make(map[string]Client, len(b.clients))
-
-	for id, c := range b.clients {
-		copy[id] = c
-	}
-
-	return copy
+// ChunkedClients splits the current clients into a chunked slice.
+func (b *Broker) ChunkedClients(chunkSize int) [][]Client {
+	return list.ToChunks(b.store.Values(), chunkSize)
 }
 
 // ClientById finds a registered client by its id.
 //
 // Returns non-nil error when client with clientId is not registered.
 func (b *Broker) ClientById(clientId string) (Client, error) {
-	b.mux.RLock()
-	defer b.mux.RUnlock()
-
-	client, ok := b.clients[clientId]
+	client, ok := b.store.GetOk(clientId)
 	if !ok {
-		return nil, fmt.Errorf("No client associated with connection ID %q", clientId)
+		return nil, fmt.Errorf("no client associated with connection ID %q", clientId)
 	}
 
 	return client, nil
@@ -50,21 +44,17 @@ func (b *Broker) ClientById(clientId string) (Client, error) {
 
 // Register adds a new client to the broker instance.
 func (b *Broker) Register(client Client) {
-	b.mux.Lock()
-	defer b.mux.Unlock()
-
-	b.clients[client.Id()] = client
+	b.store.Set(client.Id(), client)
 }
 
 // Unregister removes a single client by its id.
 //
 // If client with clientId doesn't exist, this method does nothing.
 func (b *Broker) Unregister(clientId string) {
-	b.mux.Lock()
-	defer b.mux.Unlock()
-
-	if client, ok := b.clients[clientId]; ok {
-		client.Discard()
-		delete(b.clients, clientId)
+	client := b.store.Get(clientId)
+	if client == nil {
+		return
 	}
+	client.Discard()
+	b.store.Remove(clientId)
 }

tools/subscriptions/broker_test.go

@@ -36,6 +36,32 @@ func TestClients(t *testing.T) {
 	}
 }
 
+func TestChunkedClients(t *testing.T) {
+	b := subscriptions.NewBroker()
+
+	chunks := b.ChunkedClients(2)
+	if total := len(chunks); total != 0 {
+		t.Fatalf("Expected %d chunks, got %d", 0, total)
+	}
+
+	b.Register(subscriptions.NewDefaultClient())
+	b.Register(subscriptions.NewDefaultClient())
+	b.Register(subscriptions.NewDefaultClient())
+
+	chunks = b.ChunkedClients(2)
+	if total := len(chunks); total != 2 {
+		t.Fatalf("Expected %d chunks, got %d", 2, total)
+	}
+
+	if total := len(chunks[0]); total != 2 {
+		t.Fatalf("Expected the first chunk to have 2 clients, got %d", total)
+	}
+
+	if total := len(chunks[1]); total != 1 {
+		t.Fatalf("Expected the second chunk to have 1 client, got %d", total)
+	}
+}
+
 func TestClientById(t *testing.T) {
 	b := subscriptions.NewBroker()
 

tools/subscriptions/client.go

@@ -22,11 +22,8 @@ type Message struct {
 // SubscriptionOptions defines the request options (query params, headers, etc.)
 // for a single subscription topic.
 type SubscriptionOptions struct {
-	// @todo after the requests handling refactoring consider
-	// changing to map[string]string or map[string][]string
-
-	Query   map[string]any `json:"query"`
-	Headers map[string]any `json:"headers"`
+	Query   map[string]string `json:"query"`
+	Headers map[string]string `json:"headers"`
 }
 
 // Client is an interface for a generic subscription client.
@@ -168,25 +165,33 @@ func (c *DefaultClient) Subscribe(subs ...string) {
 		}
 
 		// extract subscription options (if any)
-		options := SubscriptionOptions{}
+		rawOptions := struct {
+			// note: any instead of string to minimize the breaking changes with earlier versions
+			Query   map[string]any `json:"query"`
+			Headers map[string]any `json:"headers"`
+		}{}
 		u, err := url.Parse(s)
 		if err == nil {
-			rawOptions := u.Query().Get(optionsParam)
-			if rawOptions != "" {
-				json.Unmarshal([]byte(rawOptions), &options)
+			raw := u.Query().Get(optionsParam)
+			if raw != "" {
+				json.Unmarshal([]byte(raw), &rawOptions)
 			}
 		}
 
+		options := SubscriptionOptions{
+			Query:   make(map[string]string, len(rawOptions.Query)),
+			Headers: make(map[string]string, len(rawOptions.Headers)),
+		}
+
 		// normalize query
 		// (currently only single string values are supported for consistency with the default routes handling)
-		for k, v := range options.Query {
+		for k, v := range rawOptions.Query {
 			options.Query[k] = cast.ToString(v)
 		}
 
 		// normalize headers name and values, eg. "X-Token" is converted to "x_token"
 		// (currently only single string values are supported for consistency with the default routes handling)
-		for k, v := range options.Headers {
-			delete(options.Headers, k)
+		for k, v := range rawOptions.Headers {
 			options.Headers[inflector.Snakecase(k)] = cast.ToString(v)
 		}
 

tools/subscriptions/client_test.go

@@ -126,7 +126,7 @@ func TestSubscribeOptions(t *testing.T) {
 		name            string
 		expectedOptions string
 	}{
-		{sub1, `{"query":null,"headers":null}`},
+		{sub1, `{"query":{},"headers":{}}`},
 		{sub2, `{"query":{"name":"123"},"headers":{"x_token":"456"}}`},
 	}
 
@@ -144,7 +144,7 @@ func TestSubscribeOptions(t *testing.T) {
 			rawStr := string(rawBytes)
 
 			if rawStr != s.expectedOptions {
-				t.Fatalf("Expected options \n%v \ngot \n%v", s.expectedOptions, rawStr)
+				t.Fatalf("Expected options\n%v\ngot\n%v", s.expectedOptions, rawStr)
 			}
 		})
 	}