add context timeout on cancellation

pkg/codegen/generators/templates/controller.tmpl

@@ -186,7 +186,8 @@ func (c *{{ $.Prefix }}Controller) Subscribe{{namify $key}}(ctx context.Context,
     return nil
 }
 
-// Unsubscribe{{namify $key}} will unsubscribe messages from '{{$key}}' channel
+// Unsubscribe{{namify $key}} will unsubscribe messages from '{{$key}}' channel.
+// A timeout can be set in context to avoid blocking operation, if needed.
 {{- if .Parameters}}
 func (c *{{ $.Prefix }}Controller) Unsubscribe{{namify $key}}(ctx context.Context, params {{namify $key}}Parameters) {
 {{- else}}
@@ -205,7 +206,7 @@ func (c *{{ $.Prefix }}Controller) Unsubscribe{{namify $key}}(ctx context.Contex
     ctx = add{{ $.Prefix }}ContextValues(ctx, path)
 
     // Stop the subscription
-    sub.Cancel()
+    sub.Cancel(ctx)
 
     // Remove if from the subscribers
     delete(c.subscriptions, path)
@@ -261,10 +262,12 @@ func (c *{{ $.Prefix }}Controller) Publish{{namify $key}}(ctx context.Context, m
 {{if eq .Prefix "User" -}}
 {{- range  $key, $value := .SubscribeChannels -}}
 {{- if ne $value.Subscribe.Message.CorrelationIDLocation ""}}
-// WaitFor{{namify $key}} will wait for a specific message by its correlation ID
+// WaitFor{{namify $key}} will wait for a specific message by its correlation ID.
 //
-// The pub function is the publication function that should be used to send the message
-// It will be called after subscribing to the channel to avoid race condition, and potentially loose the message
+// The pub function is the publication function that should be used to send the message.
+// It will be called after subscribing to the channel to avoid race condition, and potentially loose the message.
+//
+// A timeout can be set in context to avoid blocking operation, if needed.
 {{- if .Parameters}}
 func (cc *UserController) WaitFor{{namify $key}}(ctx context.Context, params {{namify $key}}Parameters, publishMsg MessageWithCorrelationID, pub func(ctx context.Context) error) ({{channelToMessageTypeName $value}}, error) {
 {{- else}}
@@ -287,7 +290,7 @@ func (cc *UserController) WaitFor{{namify $key}}(ctx context.Context, publishMsg
     // Close subscriber on leave
     defer func(){
         // Stop the subscription
-        sub.Cancel()
+        sub.Cancel(ctx)
 
         // Logging unsubscribing
         cc.logger.Info(ctx, "Unsubscribed from channel")

pkg/extensions/broker.go

@@ -53,13 +53,14 @@ func (bcs BrokerChannelSubscription) WaitForCancellationAsync(cleanup func()) {
 // Cancel cancels the subscription from user perspective. It will ask for clean
 // up on broker, which will return when finished to avoid dangling resources, such
 // as non-existent queue listeners on (broker) server side.
-func (bcs BrokerChannelSubscription) Cancel() {
+func (bcs BrokerChannelSubscription) Cancel(ctx context.Context) {
 	// Send a cancellation request
 	bcs.cancel <- true
 
 	// Wait for the cancellation to be effective
-	for open := true; open; {
-		_, open = <-bcs.cancel
+	select {
+	case <-bcs.cancel:
+	case <-ctx.Done():
 	}
 }
 

pkg/extensions/versioning/broker_subscription.go

@@ -30,7 +30,7 @@ func newBrokerSubscription(
 	}
 }
 
-func (bs *brokerSubscription) createVersionListener(version string) (versionSubcription, error) {
+func (bs *brokerSubscription) createVersionListener(ctx context.Context, version string) (versionSubcription, error) {
 	// Lock the versions to avoid conflict
 	bs.versionsMutex.Lock()
 	defer bs.versionsMutex.Unlock()
@@ -44,12 +44,12 @@ func (bs *brokerSubscription) createVersionListener(version string) (versionSubc
 	// Create the channels necessary
 	cbv := newVersionSubscription(version, bs)
 	bs.versionsChannels[version] = cbv
-	defer cbv.launchListener()
+	defer cbv.launchListener(ctx)
 
 	return cbv, nil
 }
 
-func (bs *brokerSubscription) removeVersionListener(vs *versionSubcription) {
+func (bs *brokerSubscription) removeVersionListener(ctx context.Context, vs *versionSubcription) {
 	// Lock the versions to avoid conflict
 	bs.versionsMutex.Lock()
 	defer bs.versionsMutex.Unlock()
@@ -67,7 +67,7 @@ func (bs *brokerSubscription) removeVersionListener(vs *versionSubcription) {
 	}
 
 	// Otherwise cancel the broker listener
-	bs.subscription.Cancel()
+	bs.subscription.Cancel(ctx)
 
 	// Then delete the channelsByBroker from the Version Switch Wrapper
 	delete(bs.parent.channels, bs.channel)

pkg/extensions/versioning/version_subscription.go

@@ -1,6 +1,9 @@
 package versioning
 
 import (
+	"context"
+	"time"
+
 	"github.com/lerenn/asyncapi-codegen/pkg/extensions"
 	"github.com/lerenn/asyncapi-codegen/pkg/extensions/brokers"
 )
@@ -22,9 +25,14 @@ func newVersionSubscription(version string, parent *brokerSubscription) versionS
 	}
 }
 
-func (vs *versionSubcription) launchListener() {
+func (vs *versionSubcription) launchListener(ctx context.Context) {
 	// Wait for cancellation and remove version listener when it happens
 	vs.subscription.WaitForCancellationAsync(func() {
-		vs.parent.removeVersionListener(vs)
+		// Create cancel function in case there is a problem with broker removal
+		ctx, cancel := context.WithTimeout(ctx, 10*time.Second)
+		defer cancel()
+
+		// Remove the version listener
+		vs.parent.removeVersionListener(ctx, vs)
 	})
 }

pkg/extensions/versioning/wrapper.go

@@ -115,7 +115,7 @@ func (w *Wrapper) Subscribe(ctx context.Context, channel string) (extensions.Bro
 	}
 
 	// Check if the version already exists
-	cbv, err := brokerChannel.createVersionListener(version)
+	cbv, err := brokerChannel.createVersionListener(ctx, version)
 
 	return cbv.subscription, err
 }