Repository: ossf/package-analysis (language: python)

Issue:
PubSub messages being redelivered while a worker is actively processing the message.
GCP PubSub has a 600 second (10m) window for messages to be Acked.

If the message is not acked in this period it will be scheduled for redelivery.

The workers in Package Analysis can sometimes take over 10m to run, which causes multiple workers to be scheduled to process the same message - effectively starving the cluster of resources over time.

To avoid this the window needs to be greater than 10m, but GCP does not allow values higher than this.

A possible work around is to issue `ModifyAckDeadline` messages to GCP to extend the message's deadline while it is still working on it. Ideally [`time.Ticker`](https://pkg.go.dev/time#Ticker) can be used to deliver updates periodically until the analysis is complete.


Discussion / hints:

Fix patch:
diff --git a/cmd/worker/main.go b/cmd/worker/main.go
index 06301110..a8557369 100644
--- a/cmd/worker/main.go
+++ b/cmd/worker/main.go
@@ -19,6 +19,7 @@ import (
 	_ "gocloud.dev/pubsub/gcppubsub"
 	_ "gocloud.dev/pubsub/kafkapubsub"
 
+	"github.com/ossf/package-analysis/cmd/worker/pubsubextender"
 	"github.com/ossf/package-analysis/internal/featureflags"
 	"github.com/ossf/package-analysis/internal/log"
 	"github.com/ossf/package-analysis/internal/notification"
@@ -96,7 +97,6 @@ func handleMessage(ctx context.Context, msg *pubsub.Message, packagesBucket *blo
 	name := msg.Metadata["name"]
 	if name == "" {
 		log.Warn("name is empty")
-		msg.Ack()
 		return nil
 	}
 
@@ -104,7 +104,6 @@ func handleMessage(ctx context.Context, msg *pubsub.Message, packagesBucket *blo
 	if ecosystem == "" {
 		log.Warn("ecosystem is empty",
 			"name", name)
-		msg.Ack()
 		return nil
 	}
 
@@ -113,7 +112,6 @@ func handleMessage(ctx context.Context, msg *pubsub.Message, packagesBucket *blo
 		log.Warn("Unsupported pkg manager",
 			log.Label("ecosystem", ecosystem.String()),
 			"name", name)
-		msg.Ack()
 		return nil
 	}
 
@@ -187,7 +185,6 @@ func handleMessage(ctx context.Context, msg *pubsub.Message, packagesBucket *blo
 		}
 	}
 
-	msg.Ack()
 	return nil
 }
 
@@ -196,6 +193,11 @@ func messageLoop(ctx context.Context, subURL, packagesBucket, notificationTopicU
 	if err != nil {
 		return err
 	}
+	extender, err := pubsubextender.New(ctx, subURL, sub)
+	if err != nil {
+		return err
+	}
+	log.Info("Subscription deadline extender", "deadline", extender.Deadline, "grace_period", extender.GracePeriod)
 
 	// the default value of the notificationTopic object is nil
 	// if no environment variable for a notification topic is set,
@@ -227,10 +229,31 @@ func messageLoop(ctx context.Context, subURL, packagesBucket, notificationTopicU
 			// All subsequent receive calls will return the same error, so we bail out.
 			return fmt.Errorf("error receiving message: %w", err)
 		}
+		me, err := extender.Start(ctx, msg, func() {
+			log.Info("Message Ack deadline extended", "message_id", msg.LoggableID, "message_meta", msg.Metadata)
+		})
+		if err != nil {
+			// If Start fails it will always fail, so we bail out.
+			// Nack the message if we can to indicate the failure.
+			if msg.Nackable() {
+				msg.Nack()
+			}
+			return fmt.Errorf("error starting message ack deadline extender: %w", err)
+		}
 
 		if err := handleMessage(ctx, msg, pkgsBkt, resultsBuckets, imageSpec, notificationTopic); err != nil {
-			log.Error("Failed to process message",
-				"error", err)
+			log.Error("Failed to process message", "error", err)
+			if err := me.Stop(); err != nil {
+				log.Error("Extender failed", "error", err)
+			}
+			if msg.Nackable() {
+				msg.Nack()
+			}
+		} else {
+			if err := me.Stop(); err != nil {
+				log.Error("Extender failed", "error", err)
+			}
+			msg.Ack()
 		}
 	}
 }
diff --git a/cmd/worker/pubsubextender/extender.go b/cmd/worker/pubsubextender/extender.go
new file mode 100644
index 00000000..87968397
--- /dev/null
+++ b/cmd/worker/pubsubextender/extender.go
@@ -0,0 +1,134 @@
+package pubsubextender
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"net/url"
+	"time"
+
+	"gocloud.dev/pubsub"
+	"gocloud.dev/pubsub/gcppubsub"
+)
+
+const (
+	defaultGracePeriod = 60 * time.Second
+	defaultDeadline    = 300 * time.Second
+)
+
+var ErrInvalidGracePeriod = errors.New("invalid grace period")
+
+type driver interface {
+	// ExtendMessageDeadline sends a request to the pubsub service to change
+	// the deadline for the supplied message to the supplied deadline duration.
+	ExtendMessageDeadline(ctx context.Context, msg *pubsub.Message, deadline time.Duration) error
+
+	// GetSubscriptionDeadline is used to get the existing deadline period for
+	// the pubsub subscription. This deadline is used by default when extending
+	// the message deadline.
+	GetSubscriptionDeadline(ctx context.Context) (time.Duration, error)
+}
+
+type Extender struct {
+	driver      driver
+	Deadline    time.Duration
+	GracePeriod time.Duration
+}
+
+func New(ctx context.Context, subURL string, sub *pubsub.Subscription) (*Extender, error) {
+	u, err := url.Parse(subURL)
+	if err != nil {
+		return nil, err
+	}
+
+	var d driver
+	switch u.Scheme {
+	case gcppubsub.Scheme:
+		d, err = newGCPDriver(u, sub)
+		if err != nil {
+			return nil, err
+		}
+	default:
+		d = &noopDriver{}
+	}
+
+	deadline, err := d.GetSubscriptionDeadline(ctx)
+	if err != nil {
+		return nil, err
+	}
+	if deadline == 0 {
+		deadline = defaultDeadline
+	}
+
+	return &Extender{
+		driver:      d,
+		Deadline:    deadline,
+		GracePeriod: defaultGracePeriod,
+	}, nil
+}
+
+type MessageExtender struct {
+	ticker   *time.Ticker
+	msg      *pubsub.Message
+	done     chan bool
+	exited   chan error
+	callback func()
+	running  bool
+}
+
+func (e *Extender) Start(ctx context.Context, msg *pubsub.Message, callback func()) (*MessageExtender, error) {
+	freq := e.Deadline - e.GracePeriod
+	if freq < 0 {
+		// GracePeriod is larger than Deadline.
+		return nil, fmt.Errorf("%w: deadline %v is smaller than grace period %v", ErrInvalidGracePeriod, e.Deadline, e.GracePeriod)
+	}
+
+	me := &MessageExtender{
+		ticker:   time.NewTicker(freq),
+		msg:      msg,
+		done:     make(chan bool),
+		exited:   make(chan error),
+		callback: callback,
+		running:  true,
+	}
+
+	go func(ctx context.Context, me *MessageExtender, e *Extender) {
+		var err error
+		for {
+			select {
+			case <-me.done:
+				me.ticker.Stop()
+				me.exited <- err
+				return
+			case <-me.ticker.C:
+				err = e.driver.ExtendMessageDeadline(ctx, me.msg, e.Deadline)
+				if err != nil {
+					me.ticker.Stop() // don't send any more ticks.
+					// at this point the goroutine waits for done so that it
+					// can exit cleanly and report the error without race
+					// conditions.
+				} else if me.callback != nil {
+					me.callback()
+				}
+			}
+		}
+	}(ctx, me, e)
+	return me, nil
+}
+
+func (me *MessageExtender) IsRunning() bool {
+	return me.running
+}
+
+func (me *MessageExtender) Stop() error {
+	if me.running {
+		// Signal the goroutine that stop has been called and we are done.
+		me.done <- true
+		// Wait for the goroutine to exit and collect any error that it may
+		// have waiting for us.
+		err := <-me.exited
+		me.running = false
+		return err
+	}
+	return nil
+}
diff --git a/cmd/worker/pubsubextender/gcpdriver.go b/cmd/worker/pubsubextender/gcpdriver.go
new file mode 100644
index 00000000..7208b1d0
--- /dev/null
+++ b/cmd/worker/pubsubextender/gcpdriver.go
@@ -0,0 +1,86 @@
+package pubsubextender
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"net/url"
+	"path"
+	"regexp"
+	"strings"
+	"time"
+
+	api "cloud.google.com/go/pubsub/apiv1"
+	pb "cloud.google.com/go/pubsub/apiv1/pubsubpb"
+	"gocloud.dev/pubsub"
+	"gocloud.dev/pubsub/gcppubsub"
+)
+
+const (
+	gcpMinAckDeadline = 10 * time.Second
+	gcpMaxAckDeadline = 600 * time.Second
+)
+
+var subscriptionPathRE = regexp.MustCompile("^projects/.+/subscriptions/.+$")
+
+type gcpDriver struct {
+	client *api.SubscriberClient
+	path   string
+}
+
+func newGCPDriver(u *url.URL, sub *pubsub.Subscription) (driver, error) {
+	d := &gcpDriver{}
+
+	// Ignore unsupported schemes
+	if u.Scheme != "" && u.Scheme != gcppubsub.Scheme {
+		return nil, errors.New("unsupported scheme")
+	}
+
+	subPath := path.Join(u.Host, u.Path)
+	if !subscriptionPathRE.MatchString(subPath) {
+		// assume the Host is Project ID and Path is the subscription
+		subPath = fmt.Sprintf("projects/%s/subscriptions/%s", u.Host, strings.TrimPrefix(u.Path, "/"))
+	}
+
+	var c *api.SubscriberClient
+	if !sub.As(&c) {
+		return nil, errors.New("not a GCP subscription")
+	}
+	d.client = c
+	d.path = subPath
+	return d, nil
+}
+
+// ExtendMessageDeadline implements the driver interface.
+func (d *gcpDriver) ExtendMessageDeadline(ctx context.Context, msg *pubsub.Message, deadline time.Duration) error {
+	// Ensure the deadline is within acceptable bounds.
+	if deadline < gcpMinAckDeadline {
+		deadline = gcpMinAckDeadline
+	} else if deadline > gcpMaxAckDeadline {
+		deadline = gcpMaxAckDeadline
+	}
+
+	var rm *pb.ReceivedMessage
+	if !msg.As(&rm) {
+		return errors.New("not a gcp message")
+	}
+
+	if err := d.client.ModifyAckDeadline(ctx, &pb.ModifyAckDeadlineRequest{
+		Subscription:       d.path,
+		AckIds:             []string{rm.AckId},
+		AckDeadlineSeconds: int32(deadline / time.Second),
+	}); err != nil {
+		return fmt.Errorf("failed to extend message deadline: %w", err)
+	}
+
+	return nil
+}
+
+// GetSubscriptionDeadline implements the driver interface.
+func (d *gcpDriver) GetSubscriptionDeadline(ctx context.Context) (time.Duration, error) {
+	resp, err := d.client.GetSubscription(ctx, &pb.GetSubscriptionRequest{Subscription: d.path})
+	if err != nil {
+		return 0, err
+	}
+	return time.Duration(resp.GetAckDeadlineSeconds()), nil
+}
diff --git a/cmd/worker/pubsubextender/noopdriver.go b/cmd/worker/pubsubextender/noopdriver.go
new file mode 100644
index 00000000..2b79fce5
--- /dev/null
+++ b/cmd/worker/pubsubextender/noopdriver.go
@@ -0,0 +1,20 @@
+package pubsubextender
+
+import (
+	"context"
+	"time"
+
+	"gocloud.dev/pubsub"
+)
+
+type noopDriver struct{}
+
+// ExtendMessageDeadline implements the driver interface.
+func (d *noopDriver) ExtendMessageDeadline(ctx context.Context, msg *pubsub.Message, deadline time.Duration) error {
+	return nil
+}
+
+// GetSubscriptionDeadline implements the driver interface.
+func (d *noopDriver) GetSubscriptionDeadline(ctx context.Context) (time.Duration, error) {
+	return 0, nil
+}


Test patch:
diff --git a/cmd/worker/pubsubextender/extender_test.go b/cmd/worker/pubsubextender/extender_test.go
new file mode 100644
index 00000000..28139e1d
--- /dev/null
+++ b/cmd/worker/pubsubextender/extender_test.go
@@ -0,0 +1,139 @@
+package pubsubextender
+
+import (
+	"context"
+	"errors"
+	"testing"
+	"time"
+
+	"gocloud.dev/pubsub"
+)
+
+type testDriver struct {
+	extendCount int
+	lastExtend  struct {
+		msg      *pubsub.Message
+		deadline time.Duration
+	}
+	err error
+}
+
+// ExtendMessageDeadline implements the driver interface.
+func (d *testDriver) ExtendMessageDeadline(ctx context.Context, msg *pubsub.Message, deadline time.Duration) error {
+	d.extendCount++
+	d.lastExtend.msg = msg
+	d.lastExtend.deadline = deadline
+	return d.err
+}
+
+// GetSubscriptionDeadline implements the driver interface.
+func (d *testDriver) GetSubscriptionDeadline(ctx context.Context) (time.Duration, error) {
+	return 0, nil
+}
+
+func TestNew(t *testing.T) {
+	e, err := New(context.Background(), "not://a/real/pubsub/subscription", nil)
+	if err != nil {
+		t.Fatalf("New() = %v; want no error", err)
+	}
+	if e.Deadline != defaultDeadline {
+		t.Errorf("Deadline = %v; want %v", e.Deadline, defaultDeadline)
+	}
+	if e.GracePeriod != defaultGracePeriod {
+		t.Errorf("GracePeriod = %v; want %v", e.GracePeriod, defaultGracePeriod)
+	}
+}
+
+func TestExtender(t *testing.T) {
+	wantDeadline := 123 * time.Millisecond
+	d := &testDriver{}
+	e := &Extender{
+		driver:      d,
+		Deadline:    wantDeadline,
+		GracePeriod: 50 * time.Millisecond,
+	}
+	ctx := context.Background()
+	wantMsg := &pubsub.Message{
+		LoggableID: "test-msg-0",
+		Metadata:   map[string]string{"test": "test"},
+	}
+
+	me, err := e.Start(ctx, wantMsg, func() {
+		if got := d.lastExtend.msg; got != wantMsg {
+			t.Errorf("ExtendMessageDeadline got message %v, want %v", got, wantMsg)
+		}
+		if got := d.lastExtend.deadline; got != wantDeadline {
+			t.Errorf("ExtendMessageDeadline got deadline %v, want %v", got, wantDeadline)
+		}
+
+		d.lastExtend.msg = nil
+		d.lastExtend.deadline = 0
+	})
+	if err != nil {
+		t.Fatalf("Start() = %v; want no error", err)
+	}
+	if !me.IsRunning() {
+		t.Errorf("IsRunning()#1 = false; want true")
+	}
+
+	time.Sleep(500 * time.Millisecond)
+
+	if err := me.Stop(); err != nil {
+		t.Errorf("Err() = %v; want nil", err)
+	}
+
+	if d.extendCount == 0 {
+		t.Errorf("ExtendMessageDeadline never called")
+	}
+	if me.IsRunning() {
+		t.Errorf("IsRunning()#2 = true; want false")
+	}
+
+	// Reset so we can ensure that it has stoppped.
+	d.extendCount = 0
+
+	time.Sleep(500 * time.Millisecond)
+	if d.extendCount != 0 {
+		t.Errorf("Extender not stopped")
+	}
+
+	// Calling stop again does nothing and has no error.
+	if err := me.Stop(); err != nil {
+		t.Errorf("Err() = %v; want nil", err)
+	}
+}
+
+func TestExtender_Error(t *testing.T) {
+	wantErr := errors.New("failed")
+	d := &testDriver{
+		err: wantErr,
+	}
+	e := &Extender{
+		driver:      d,
+		Deadline:    100 * time.Millisecond,
+		GracePeriod: 50 * time.Millisecond,
+	}
+	ctx := context.Background()
+	wantMsg := &pubsub.Message{
+		LoggableID: "test-msg-0",
+		Metadata:   map[string]string{"test": "test"},
+	}
+
+	me, err := e.Start(ctx, wantMsg, nil)
+	if err != nil {
+		t.Fatalf("Start() = %v; want no error", err)
+	}
+
+	time.Sleep(500 * time.Millisecond)
+
+	if err := me.Stop(); err == nil {
+		t.Errorf("Err() = %v; want %v", err, wantErr)
+	}
+	if d.extendCount != 1 {
+		t.Errorf("ExtendMessageDeadline called after error")
+	}
+	// Calling stop again does nothing and has no error.
+	if err := me.Stop(); err != nil {
+		t.Errorf("Err() = %v; want nil", err)
+	}
+}