Run subscriptions when app starts rather than app start (#346)

Author: Umang01-hash

pkg/gofr/container/container.go

@@ -30,7 +30,7 @@ type Container struct {
 
 	Services       map[string]service.HTTP
 	metricsManager metrics.Manager
-	pubsub         pubsub.Client
+	PubSub         pubsub.Client
 
 	Redis *redis.Redis
 	SQL   *sql.DB
@@ -82,15 +82,15 @@ func (c *Container) Create(conf config.Config) {
 			partition, _ := strconv.Atoi(conf.GetOrDefault("PARTITION_SIZE", "0"))
 			offSet, _ := strconv.Atoi(conf.GetOrDefault("PUBSUB_OFFSET", "-1"))
 
-			c.pubsub = kafka.New(kafka.Config{
+			c.PubSub = kafka.New(kafka.Config{
 				Broker:          conf.Get("PUBSUB_BROKER"),
 				Partition:       partition,
 				ConsumerGroupID: conf.Get("CONSUMER_ID"),
 				OffSet:          offSet,
 			}, c.Logger, c.metricsManager)
 		}
 	case "GOOGLE":
-		c.pubsub = google.New(google.Config{
+		c.PubSub = google.New(google.Config{
 			ProjectID:        conf.Get("GOOGLE_PROJECT_ID"),
 			SubscriptionName: conf.Get("GOOGLE_SUBSCRIPTION_NAME"),
 		}, c.Logger, c.metricsManager)
@@ -146,9 +146,9 @@ func (c *Container) GetAppVersion() string {
 }
 
 func (c *Container) GetPublisher() pubsub.Publisher {
-	return c.pubsub
+	return c.PubSub
 }
 
 func (c *Container) GetSubscriber() pubsub.Subscriber {
-	return c.pubsub
+	return c.PubSub
 }

pkg/gofr/container/container_test.go

@@ -58,7 +58,7 @@ func Test_newContainerPubSubIntializationFail(t *testing.T) {
 	for _, tc := range testCases {
 		c := NewContainer(testutil.NewMockConfig(tc.configs))
 
-		assert.Nil(t, c.pubsub)
+		assert.Nil(t, c.PubSub)
 	}
 }
 
@@ -112,7 +112,7 @@ func TestContainer_GetAppVersion(t *testing.T) {
 func TestContainer_GetPublisher(t *testing.T) {
 	publisher := &mockPubSub{}
 
-	c := &Container{pubsub: publisher}
+	c := &Container{PubSub: publisher}
 
 	out := c.GetPublisher()
 
@@ -122,7 +122,7 @@ func TestContainer_GetPublisher(t *testing.T) {
 func TestContainer_GetSubscriber(t *testing.T) {
 	subscriber := &mockPubSub{}
 
-	c := &Container{pubsub: subscriber}
+	c := &Container{PubSub: subscriber}
 
 	out := c.GetSubscriber()
 

pkg/gofr/gofr.go

@@ -1,7 +1,6 @@
 package gofr
 
 import (
-	"context"
 	"fmt"
 
 	"net/http"
@@ -39,9 +38,10 @@ type App struct {
 	// container is unexported because this is an internal implementation and applications are provided access to it via Context
 	container *container.Container
 
-	grpcRegistered       bool
-	httpRegistered       bool
-	subscriberRegistered bool
+	grpcRegistered bool
+	httpRegistered bool
+
+	subscriptionManager SubscriptionManager
 }
 
 // RegisterService adds a grpc service to the gofr application.
@@ -83,6 +83,8 @@ func New() *App {
 
 	app.grpcServer = newGRPCServer(app.container, port)
 
+	app.subscriptionManager = newSubscriptionManager(app.container)
+
 	return app
 }
 
@@ -148,7 +150,12 @@ func (a *App) Run() {
 	}
 
 	// If subscriber is registered, block main go routine to wait for subscriber to receive messages
-	if a.subscriberRegistered {
+	if a.subscriptionManager.subscriptions != nil {
+		// Start subscribers concurrently using go-routines
+		for topic, handler := range a.subscriptionManager.subscriptions {
+			go a.subscriptionManager.startSubscriber(topic, handler)
+		}
+
 		wg.Add(1)
 	}
 
@@ -264,31 +271,5 @@ func (a *App) Subscribe(topic string, handler SubscribeFunc) {
 		return
 	}
 
-	a.subscriberRegistered = true
-
-	// continuously subscribe in an infinite loop
-	go func() {
-		for {
-			msg, err := a.container.GetSubscriber().Subscribe(context.Background(), topic)
-			if msg == nil {
-				continue
-			}
-
-			if err != nil {
-				a.container.Logger.Errorf("error while reading from Kafka, err: %v", err.Error())
-
-				continue
-			}
-
-			// create a gofr context with message as request
-			ctx := newContext(nil, msg, a.container)
-
-			err = handler(ctx)
-
-			// commit the message if the subscription function does not return error
-			if err == nil {
-				msg.Commit()
-			}
-		}
-	}()
+	a.subscriptionManager.subscriptions[topic] = handler
 }

pkg/gofr/logging/level_test.go

@@ -1,6 +1,7 @@
 package logging
 
 import (
+	"os"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -93,3 +94,16 @@ func TestGetLevelFromString(t *testing.T) {
 		assert.Equal(t, tc.expected, actual, "TEST[%d], Failed.\n%s", i, tc.desc)
 	}
 }
+
+func Test_changeLevel(t *testing.T) {
+	l := logger{
+		level:      INFO,
+		normalOut:  os.Stdout,
+		errorOut:   os.Stderr,
+		isTerminal: false,
+	}
+
+	l.changeLevel(ERROR)
+
+	assert.Equal(t, ERROR, l.level, "Test_changeLevel failed! expected level to be error ")
+}

pkg/gofr/subscriber.go

@@ -1,3 +1,46 @@
 package gofr
 
+import (
+	"context"
+
+	"gofr.dev/pkg/gofr/container"
+)
+
 type SubscribeFunc func(c *Context) error
+
+type SubscriptionManager struct {
+	container     *container.Container
+	subscriptions map[string]SubscribeFunc
+}
+
+func newSubscriptionManager(c *container.Container) SubscriptionManager {
+	return SubscriptionManager{
+		container:     c,
+		subscriptions: make(map[string]SubscribeFunc),
+	}
+}
+
+func (s *SubscriptionManager) startSubscriber(topic string, handler SubscribeFunc) {
+	// continuously subscribe in an infinite loop
+	for {
+		msg, err := s.container.GetSubscriber().Subscribe(context.Background(), topic)
+		if msg == nil {
+			continue
+		}
+
+		if err != nil {
+			s.container.Logger.Errorf("error while reading from Kafka, err: %v", err.Error())
+			continue
+		}
+
+		ctx := newContext(nil, msg, s.container)
+		err = handler(ctx)
+
+		// commit the message if the subscription function does not return error
+		if err == nil {
+			msg.Commit()
+		} else {
+			s.container.Logger.Errorf("error in handler for topic %s: %v", topic, err)
+		}
+	}
+}

pkg/gofr/subscriber_test.go

@@ -0,0 +1,114 @@
+package gofr
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"strings"
+	"testing"
+	"time"
+
+	"gofr.dev/pkg/gofr/container"
+	"gofr.dev/pkg/gofr/datasource/pubsub"
+	"gofr.dev/pkg/gofr/logging"
+	"gofr.dev/pkg/gofr/testutil"
+)
+
+var errHandler = errors.New("error in subscribing")
+
+func handleError(err string) error {
+	return fmt.Errorf("%w: %s", errHandler, err)
+}
+
+var errSubscription = errors.New("subscription error")
+
+func subscriptionError(err string) error {
+	return fmt.Errorf("%w: %s", errSubscription, err)
+}
+
+type mockSubscriber struct {
+}
+
+func (s mockSubscriber) Publish(_ context.Context, _ string, _ []byte) error {
+	return nil
+}
+
+type Message struct {
+	Topic    string
+	Value    []byte
+	MetaData interface{}
+}
+
+func (mockSubscriber) Subscribe(_ context.Context, topic string) (*pubsub.Message, error) {
+	if topic == "test-topic" {
+		return &pubsub.Message{
+			Topic: "test-topic",
+			Value: []byte(`{"data":{"productId":"123","price":"599"}}`),
+		}, nil
+	}
+
+	return &pubsub.Message{
+		Topic: "test-topic",
+		Value: []byte(`{"data":{"productId":"123","price":"599"}}`),
+	}, subscriptionError("subscription error")
+}
+
+func TestSubscriptionManager_HandlerError(t *testing.T) {
+	done := make(chan struct{})
+
+	testLogs := testutil.StderrOutputForFunc(func() {
+		mockContainer := container.Container{
+			Logger: logging.NewLogger(logging.ERROR),
+			PubSub: mockSubscriber{},
+		}
+		subscriptionManager := newSubscriptionManager(&mockContainer)
+
+		// Run the subscriber in a goroutine
+		go func() {
+			subscriptionManager.startSubscriber("test-topic",
+				func(c *Context) error {
+					return handleError("error in test-topic")
+				})
+		}()
+
+		// this sleep is added to wait for StderrOutputForFunc to collect the logs inside the testLogs
+		time.Sleep(1 * time.Millisecond)
+	})
+
+	// signal the test to end
+	close(done)
+
+	if !strings.Contains(testLogs, "error in handler for topic test-topic: error in subscribing: error in test-topic") {
+		t.Error("TestSubscriptionManager_HandlerError Failed! Missing log message about handler error")
+	}
+}
+
+func TestSubscriptionManager_SubscribeError(t *testing.T) {
+	done := make(chan struct{})
+
+	testLogs := testutil.StderrOutputForFunc(func() {
+		mockContainer := container.Container{
+			Logger: logging.NewLogger(logging.ERROR),
+			PubSub: mockSubscriber{},
+		}
+		subscriptionManager := newSubscriptionManager(&mockContainer)
+
+		// Run the subscriber in a goroutine
+		go func() {
+			subscriptionManager.startSubscriber("abc",
+				func(c *Context) error {
+					return handleError("error in abc")
+				})
+		}()
+
+		// this sleep is added to wait for StderrOutputForFunc to collect the logs inside the testLogs
+		time.Sleep(1 * time.Millisecond)
+	})
+
+	// signal the test to end
+	close(done)
+
+	if !strings.Contains(testLogs, "error while reading from Kafka, err: ") {
+		t.Error("TestSubscriptionManager_SubscribeError Failed! Missing log message about subscription error")
+	}
+}

pkg/gofr/testutil/mock_logger.go

@@ -20,6 +20,14 @@ type MockLogger struct {
 	out    io.Writer
 	errOut io.Writer
 }
+type (
+	Level int
+)
+
+//nolint:unused //changeLevel is present to satisfy the mocking of logger inside testutil.
+func (m *MockLogger) changeLevel(level Level) {
+	m.level = int(level)
+}
 
 func NewMockLogger(level int) *MockLogger {
 	return &MockLogger{