fix(pubsub): proper bidirectional blocking

picosh · Sep 9, 2024 · 354f418 · 354f418
1 parent fdadd78
commit 354f418
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 5 deletions.
diff --git a/go.mod b/go.mod
@@ -36,7 +36,7 @@ require (
 	github.com/muesli/termenv v0.15.3-0.20240509142007-81b8f94111d5
 	github.com/neurosnap/go-exif-remove v0.0.0-20221010134343-50d1e3c35577
 	github.com/picosh/pobj v0.0.0-20240709135546-27097077b26a
-	github.com/picosh/pubsub v0.0.0-20240908153701-d8e78982a3da
+	github.com/picosh/pubsub v0.0.0-20240909042445-92777a8b167b
 	github.com/picosh/send v0.0.0-20240820031602-5d3b1a4494cc
 	github.com/picosh/tunkit v0.0.0-20240709033345-8315d4f3cd0e
 	github.com/sabhiram/go-gitignore v0.0.0-20210923224102-525f6e181f06

diff --git a/go.sum b/go.sum
@@ -224,8 +224,8 @@ github.com/picosh/go-rsync-receiver v0.0.0-20240709135253-1daf4b12a9fc h1:bvcsoO
 github.com/picosh/go-rsync-receiver v0.0.0-20240709135253-1daf4b12a9fc/go.mod h1:i0iR3W4GSm1PuvVxB9OH32E5jP+CYkVb2NQSe0JCtlo=
 github.com/picosh/pobj v0.0.0-20240709135546-27097077b26a h1:Cr1xODiyd/SjjBRtYA9VX6Do3D+w+DansQzkb4NGeyA=
 github.com/picosh/pobj v0.0.0-20240709135546-27097077b26a/go.mod h1:VIkR1MZBvxSK2OO47jikxikAO/sKb/vTmXX5ZuYTIvo=
-github.com/picosh/pubsub v0.0.0-20240908153701-d8e78982a3da h1:HCu4NyCoexe4Xjm05tWnx8QFxzRXxMoeixDqN6kiZG4=
-github.com/picosh/pubsub v0.0.0-20240908153701-d8e78982a3da/go.mod h1:FKC8uot+40iXmuDzTfbxYDG5PIc3ghwkmP2iItBKH0I=
+github.com/picosh/pubsub v0.0.0-20240909042445-92777a8b167b h1:/gGhT8y9rnrv8K9ZJKZYzdWvZcnazl8NGE1DGNrD8HU=
+github.com/picosh/pubsub v0.0.0-20240909042445-92777a8b167b/go.mod h1:FKC8uot+40iXmuDzTfbxYDG5PIc3ghwkmP2iItBKH0I=
 github.com/picosh/send v0.0.0-20240820031602-5d3b1a4494cc h1:IIsJuAFG2ju3cygKVKTIsYYZf21q5S3Dr1H4fGbfgJg=
 github.com/picosh/send v0.0.0-20240820031602-5d3b1a4494cc/go.mod h1:RAgLDK3LrDK6pNeXtU9tjo28obl5DxShcTUk2nm/KCM=
 github.com/picosh/senpai v0.0.0-20240503200611-af89e73973b0 h1:pBRIbiCj7K6rGELijb//dYhyCo8A3fvxW5dijrJVtjs=

diff --git a/pubsub/cli.go b/pubsub/cli.go
@@ -103,6 +103,7 @@ func WishMiddleware(handler *CliHandler) wish.Middleware {
 
 	return func(next ssh.Handler) ssh.Handler {
 		return func(sesh ssh.Session) {
+			ctx := sesh.Context()
 			user, err := getUser(sesh, dbpool)
 			if err != nil {
 				utils.ErrorHandler(sesh, err)
@@ -199,6 +200,14 @@ func WishMiddleware(handler *CliHandler) wish.Middleware {
 					wish.Println(sesh, "no subs found ... waiting")
 				}
 
+				go func() {
+					<-ctx.Done()
+					err := pubsub.PubSub.UnPub(msg)
+					if err != nil {
+						wish.Errorln(sesh, err)
+					}
+				}()
+
 				err = pubsub.PubSub.Pub(msg)
 				wish.Println(sesh, "msg sent!")
 				if err != nil {
@@ -217,12 +226,21 @@ func WishMiddleware(handler *CliHandler) wish.Middleware {
 					name = toPublicChannel(channelName)
 				}
 
-				err = pubsub.PubSub.Sub(&psub.Subscriber{
+				sub := &psub.Subscriber{
 					ID:     uuid.NewString(),
 					Name:   name,
 					Writer: sesh,
 					Chan:   make(chan error),
-				})
+				}
+
+				go func() {
+					<-ctx.Done()
+					err := pubsub.PubSub.UnSub(sub)
+					if err != nil {
+						wish.Errorln(sesh, err)
+					}
+				}()
+				err = pubsub.PubSub.Sub(sub)
 				if err != nil {
 					wish.Errorln(sesh, err)
 				}