fix: handle timeout in github

Author: Arta Asadi

discovery/pkg/worker/service.go

@@ -1,56 +1,29 @@
 package worker
 
 import (
-	"context"
-	"errors"
 	"github.com/spf13/cobra"
 	"go.uber.org/zap"
-	"os"
-	"os/signal"
-	"syscall"
-	"time"
 )
 
 func WorkerCommand() *cobra.Command {
 	cmd := &cobra.Command{
 		RunE: func(cmd *cobra.Command, args []string) error {
-			setupCtx := cmd.Context()
+			ctx := cmd.Context()
 			cmd.SilenceUsage = true
 			logger, err := zap.NewProduction()
 			if err != nil {
 				return err
 			}
 
-			sigChan := make(chan os.Signal, 1)
-			signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
-			go func() {
-				sig := <-sigChan
-				logger.Warn("Received termination signal, worker will shut down when idle timeout is reached or on next check.", zap.String("signal", sig.String()))
-			}()
-
 			w, err := NewWorker(
 				logger,
-				setupCtx,
+				cmd.Context(),
 			)
 			if err != nil {
 				return err
 			}
 
-			// Define the idle timeout duration
-			idleTimeout := 2 * time.Minute
-
-			// Run the worker with a background context, independent of signals,
-			// and pass the idle timeout.
-			logger.Info("Starting worker with idle timeout", zap.Duration("timeout", idleTimeout))
-			runErr := w.Run(context.Background(), idleTimeout) // Use context.Background()
-
-			if runErr != nil && !errors.Is(runErr, context.Canceled) { // Avoid logging error for clean context cancellations if any were added back
-				logger.Error("Worker Run exited with error", zap.Error(runErr))
-				return runErr // Propagate actual errors
-			}
-
-			logger.Info("Worker Run finished.")
-			return nil
+			return w.Run(ctx)
 		},
 	}
 

discovery/pkg/worker/worker.go

@@ -65,126 +65,49 @@ func NewWorker(
 	return w, nil
 }
 
-func (w *Worker) Run(ctx context.Context, idleTimeout time.Duration) error {
-	w.logger.Info("Starting Run loop",
-		zap.String("nats_consumer", envs.NatsConsumer),
-		zap.String("stream", envs.StreamName),
-		zap.String("topic", envs.TopicName),
-		zap.Duration("idleTimeout", idleTimeout))
-
-	idleTimer := time.NewTimer(idleTimeout)
-
-	if !idleTimer.Stop() {
-		select {
-		case <-idleTimer.C:
-		default:
+func (w *Worker) Run(ctx context.Context) error {
+	w.logger.Info("starting to consume", zap.String("url", envs.NatsURL), zap.String("consumer", envs.NatsConsumer),
+		zap.String("stream", envs.StreamName), zap.String("topic", envs.TopicName))
+
+	consumeCtx, err := w.jq.ConsumeWithConfig(ctx, envs.NatsConsumer, envs.StreamName, []string{envs.TopicName}, jetstream.ConsumerConfig{
+		Replicas:          1,
+		AckPolicy:         jetstream.AckExplicitPolicy,
+		DeliverPolicy:     jetstream.DeliverAllPolicy,
+		MaxAckPending:     -1,
+		AckWait:           time.Minute * 30,
+		InactiveThreshold: time.Hour,
+	}, []jetstream.PullConsumeOpt{
+		jetstream.PullMaxMessages(1),
+	}, func(msg jetstream.Msg) {
+		w.logger.Info("received a new job")
+
+		err := w.ProcessMessage(ctx, msg)
+		if err != nil {
+			// Log error from ProcessMessage itself (e.g., initial setup failure)
+			// Note: Errors during task.RunTask are handled within ProcessMessage's defer
+			w.logger.Error("failed during message processing setup", zap.Error(err))
 		}
-	}
-	idleTimer.Reset(idleTimeout)
-
-	var consumer jetstream.ConsumeContext
-
-	consumerCtx, cancelConsumer := context.WithCancel(ctx)
-	defer cancelConsumer()
-
-	consumerErrChan := make(chan error, 1)
-	activityChan := make(chan struct{}, 1)
-
-	go func() {
-		var consErr error
-		consumer, consErr = w.jq.ConsumeWithConfig(
-			consumerCtx,
-			envs.NatsConsumer,
-			envs.StreamName,
-			[]string{envs.TopicName},
-			jetstream.ConsumerConfig{
-				Replicas:          1,
-				AckPolicy:         jetstream.AckExplicitPolicy,
-				DeliverPolicy:     jetstream.DeliverAllPolicy,
-				MaxAckPending:     -1,
-				AckWait:           time.Minute * 30,
-				InactiveThreshold: time.Hour,
-			},
-			[]jetstream.PullConsumeOpt{
-				jetstream.PullMaxMessages(1), // Process one message at a time
-			},
-			func(msg jetstream.Msg) {
-				w.logger.Info("Received a new job")
-				err := w.ProcessMessage(consumerCtx, msg)
-				if err != nil {
-					w.logger.Error("Failed during message processing", zap.Error(err))
-				}
-
-				if ackErr := msg.Ack(); ackErr != nil {
-					w.logger.Error("Failed to send the ack message", zap.Error(ackErr))
-
-				} else {
-					w.logger.Info("Message Ack'd successfully.")
-				}
-
-				w.logger.Info("Processing a job completed, signaling activity.")
-
-				select {
-				case activityChan <- struct{}{}:
-				default:
-					w.logger.Warn("Activity channel buffer full?")
-				}
-			})
 
-		consumerErrChan <- consErr
-
-		<-consumerCtx.Done()
-		w.logger.Info("Consumer goroutine exiting.")
-
-	}()
-
-	setupErr := <-consumerErrChan
-	if setupErr != nil {
-		w.logger.Error("Failed to start NATS consumer", zap.Error(setupErr))
-		if !idleTimer.Stop() {
-			select {
-			case <-idleTimer.C:
-			default:
-			}
+		// Ack is always sent by Run after ProcessMessage finishes or fails setup
+		if ackErr := msg.Ack(); ackErr != nil {
+			w.logger.Error("failed to send the ack message", zap.Error(ackErr))
 		}
-		return setupErr
-	}
 
-	w.logger.Info("NATS consumer started successfully. Waiting for messages or idle timeout...")
+		w.logger.Info("processing a job completed")
+	})
+	if err != nil {
+		w.logger.Error("failed to start consuming messages", zap.Error(err))
+		return err
+	}
 
-	for {
-		select {
-		case <-ctx.Done():
-			w.logger.Info("Worker Run context cancelled. Initiating shutdown.")
-			cancelConsumer()
-			if consumer != nil {
-				w.logger.Info("Draining consumer due to context cancellation...")
-				consumer.Drain()
-				w.logger.Info("Consumer drained.")
-			}
-			return ctx.Err()
+	w.logger.Info("consuming messages...")
 
-		case <-activityChan:
-			w.logger.Debug("Activity detected, resetting idle timer.")
-			if !idleTimer.Stop() {
-				select {
-				case <-idleTimer.C:
-				default:
-				}
-			}
-			idleTimer.Reset(idleTimeout)
+	<-ctx.Done()
+	w.logger.Info("Main context cancelled, draining consumer...")
+	consumeCtx.Drain()
+	w.logger.Info("Consumer stopped.")
 
-		case <-idleTimer.C:
-			w.logger.Info("Idle timeout reached. Initiating shutdown.")
-			cancelConsumer()
-			if consumer != nil {
-				w.logger.Info("Draining consumer due to idle timeout...")
-				consumer.Drain()
-				w.logger.Info("Consumer drained.")
-			}
-			return nil
-		}
-	}
+	return nil
 }
 
 func (w *Worker) ProcessMessage(ctx context.Context, msg jetstream.Msg) (err error) {