dkforest

commit 7d483ca01acb778bbe3864fe647596ffde43ad4f
parent a1789f2f636386fbe56e54e5694e848411d0c3b5
Author: n0tr1v <n0tr1v@protonmail.com>
Date:   Tue, 23 May 2023 21:37:34 -0700

hack for fast graceful exit

Diffstat:
M
pkg/pubsub/pubsub.go
 | 
16
++++++++++++++++
M
pkg/web/handlers/handlers.go
 | 
6
+++++-

2 files changed, 21 insertions(+), 1 deletion(-)
diff --git a/pkg/pubsub/pubsub.go b/pkg/pubsub/pubsub.go
@@ -83,6 +83,22 @@ type Sub[T any] struct {
 	p      *PubSub[T]
 }
 
+// ReceiveTimeout2 returns a message received on the channel or timeout
+func (s *Sub[T]) ReceiveTimeout2(timeout time.Duration, c1, c2 chan bool) (topic string, msg T, err error) {
+	select {
+	case p := <-s.ch:
+		return p.topic, p.msg, nil
+	case <-time.After(timeout):
+		return topic, msg, ErrTimeout
+	case <-c1:
+		return topic, msg, ErrCancelled
+	case <-c2:
+		return topic, msg, ErrCancelled
+	case <-s.ctx.Done():
+		return topic, msg, ErrCancelled
+	}
+}
+
 // ReceiveTimeout returns a message received on the channel or timeout
 func (s *Sub[T]) ReceiveTimeout(timeout time.Duration) (topic string, msg T, err error) {
 	select {
diff --git a/pkg/web/handlers/handlers.go b/pkg/web/handlers/handlers.go
@@ -6,6 +6,7 @@ import (
 	"crypto/sha256"
 	dutils "dkforest/pkg/database/utils"
 	"dkforest/pkg/hashset"
+	"dkforest/pkg/pubsub"
 	"dkforest/pkg/utils/crypto"
 	v1 "dkforest/pkg/web/handlers/api/v1"
 	"encoding/base64"
@@ -4996,8 +4997,11 @@ Loop:
 		c.Response().Flush()
 		clrIdx++
 
-		_, msgTyp, err := sub.ReceiveTimeout(5 * time.Second)
+		_, msgTyp, err := sub.ReceiveTimeout2(5*time.Second, quit, quit1)
 		if err != nil {
+			if err == pubsub.ErrCancelled {
+				break Loop
+			}
 			continue
 		}