Merge pull request #31 from saiya/feature/multiplexer-long-polling-early-return

Storage multiplexer improvements

Author: saiya

server/domain/mock/.gitignore

@@ -0,0 +1 @@
+*.go

server/domain/storage.go

@@ -24,6 +24,8 @@ func IsStorageNonFatalError(err error) bool {
 	return errors.Is(err, ErrInvalidChannel) || errors.Is(err, ErrSubscriptionNotFound) || errors.Is(err, ErrMalformedAckHandle)
 }
 
+//go:generate mockgen -source=${GOFILE} -package=mock -destination=./mock/${GOFILE}
+
 // Storage interface is an abstraction layer of storage implementations
 type Storage interface {
 	String() string // Stringer

server/storage/multiplex/pubsub_messaging.go

@@ -2,13 +2,17 @@ package multiplex
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"sync"
+	"time"
 
 	"github.com/saiya/dsps/server/domain"
 	"github.com/saiya/dsps/server/logger"
 )
 
+const parallelFetchEarlyReturnWindow = 300 * time.Millisecond
+
 func (s *storageMultiplexer) PublishMessages(ctx context.Context, msgs []domain.Message) error {
 	_, err := s.parallelAtLeastOneSuccess(ctx, "PublishMessages", func(ctx context.Context, _ domain.StorageID, child domain.Storage) (interface{}, error) {
 		if child := child.AsPubSubStorage(); child != nil {
@@ -25,20 +29,41 @@ func (s *storageMultiplexer) FetchMessages(ctx context.Context, sl domain.Subscr
 		moreMessages bool
 		ackHandle    domain.AckHandle
 	}
-	results, err := s.parallelAtLeastOneSuccess(ctx, "FetchMessages", func(ctx context.Context, _ domain.StorageID, child domain.Storage) (interface{}, error) {
+	parallelCtx, parallelCtxCancel := context.WithCancel(ctx)
+	defer parallelCtxCancel()
+	subscriptionMissingCh := make(chan domain.StorageID, len(s.children))
+	results, err := s.parallelAtLeastOneSuccess(parallelCtx, "FetchMessages", func(ctx context.Context, storageID domain.StorageID, child domain.Storage) (interface{}, error) {
 		if child := child.AsPubSubStorage(); child != nil {
 			msgs, moreMsgs, ackHandle, err := child.FetchMessages(ctx, sl, max, waituntil)
 			if err != nil {
+				if errors.Is(err, domain.ErrSubscriptionNotFound) || errors.Is(err, domain.ErrInvalidChannel) {
+					subscriptionMissingCh <- storageID
+				}
 				return nil, err
 			}
+			if len(msgs) > 0 {
+				// If one or more storage returns messages, multiplexer should immediately return them even if other storages still polling.
+				time.AfterFunc(parallelFetchEarlyReturnWindow, parallelCtxCancel)
+			}
 			return fetchResult{msgs: msgs, moreMessages: moreMsgs, ackHandle: ackHandle}, nil
 		}
 		return nil, errMultiplexSkipped
 	})
+	close(subscriptionMissingCh)
 	if err != nil {
 		return nil, false, domain.AckHandle{}, err
 	}
 
+	for id := range subscriptionMissingCh {
+		// Subscriber missing on this storage.
+		// This situation could occur if the storage had been temporary unavailable when subscriber created.
+		// So that automatically create subscriber to receive future messages.
+		logger.Of(ctx).Debugf(`Auto-creating (recovering) subscriber %v on storage '%s' because fetch succeeded in the multiplexer but this storage reported the subscriber does not exist.`, sl, id)
+		if err := s.children[id].AsPubSubStorage().NewSubscriber(ctx, sl); err != nil {
+			logger.Of(ctx).WarnError(fmt.Sprintf("Failed to auto-create (recover) subscriber %v on storage '%s': %%w", sl, id), err)
+		}
+	}
+
 	// Note that this merge logic honors message ordering as possible.
 	// Only exception is that storages returns messages by different ordering, possible cause of such case is that client retry to publish messages.
 	// If client retried publish, no need to guarantee ordering of the messages sent concurrently with the retry.

server/storage/multiplex/storage_multiplexer_generic_test.go

@@ -0,0 +1,70 @@
+package multiplex_test
+
+import (
+	"context"
+	"fmt"
+	"testing"
+
+	"github.com/saiya/dsps/server/config"
+	"github.com/saiya/dsps/server/domain"
+	. "github.com/saiya/dsps/server/storage/multiplex"
+	"github.com/saiya/dsps/server/storage/onmemory"
+	. "github.com/saiya/dsps/server/storage/testing"
+)
+
+var onmemoryMultiplexCtor = func(onmemConfigs ...config.OnmemoryStorageConfig) StorageCtor {
+	return func(ctx context.Context, systemClock domain.SystemClock, channelProvider domain.ChannelProvider) (domain.Storage, error) {
+		storages := map[domain.StorageID]domain.Storage{}
+		for i := range onmemConfigs {
+			storage, err := onmemory.NewOnmemoryStorage(context.Background(), &(onmemConfigs[i]), systemClock, channelProvider)
+			if err != nil {
+				return nil, err
+			}
+			storages[domain.StorageID(fmt.Sprintf("storage%d", i+1))] = storage
+		}
+		return NewStorageMultiplexer(storages)
+	}
+}
+
+func TestCoreFunction(t *testing.T) {
+	CoreFunctionTest(t, onmemoryMultiplexCtor(
+		config.OnmemoryStorageConfig{
+			DisablePubSub: true,
+			DisableJwt:    true,
+		},
+		config.OnmemoryStorageConfig{
+			DisablePubSub: true,
+			DisableJwt:    true,
+		},
+	))
+}
+
+func TestPubSub(t *testing.T) {
+	PubSubTest(t, onmemoryMultiplexCtor(
+		config.OnmemoryStorageConfig{
+			DisableJwt: true,
+		},
+		config.OnmemoryStorageConfig{
+			DisablePubSub: true, // Storage without feature support
+			DisableJwt:    true,
+		},
+		config.OnmemoryStorageConfig{
+			DisableJwt: true,
+		},
+	))
+}
+
+func TestJwt(t *testing.T) {
+	JwtTest(t, onmemoryMultiplexCtor(
+		config.OnmemoryStorageConfig{
+			DisablePubSub: true,
+		},
+		config.OnmemoryStorageConfig{
+			DisablePubSub: true,
+			DisableJwt:    true, // Storage without feature support
+		},
+		config.OnmemoryStorageConfig{
+			DisablePubSub: true,
+		},
+	))
+}