Fix RetryClient context (#97)

Author: at-wat

retryclient.go

@@ -30,7 +30,7 @@ type RetryClient struct {
 	mu             sync.Mutex
 	muQueue        sync.Mutex
 	handler        Handler
-	chTask         chan func(cli Client)
+	chTask         chan func(ctx context.Context, cli Client)
 }
 
 // Handle registers the message handler.
@@ -46,23 +46,23 @@ func (c *RetryClient) Handle(handler Handler) {
 // Publish tries to publish the message and immediately return nil.
 // If it is not acknowledged to be published, the message will be queued.
 func (c *RetryClient) Publish(ctx context.Context, message *Message) error {
-	return c.pushTask(ctx, func(cli Client) {
+	return c.pushTask(ctx, func(ctx context.Context, cli Client) {
 		c.publish(ctx, false, cli, message)
 	})
 }
 
 // Subscribe tries to subscribe the topic and immediately return nil.
 // If it is not acknowledged to be subscribed, the request will be queued.
 func (c *RetryClient) Subscribe(ctx context.Context, subs ...Subscription) error {
-	return c.pushTask(ctx, func(cli Client) {
+	return c.pushTask(ctx, func(ctx context.Context, cli Client) {
 		c.subscribe(ctx, false, cli, subs...)
 	})
 }
 
 // Unsubscribe tries to unsubscribe the topic and immediately return nil.
 // If it is not acknowledged to be unsubscribed, the request will be queued.
 func (c *RetryClient) Unsubscribe(ctx context.Context, topics ...string) error {
-	return c.pushTask(ctx, func(cli Client) {
+	return c.pushTask(ctx, func(ctx context.Context, cli Client) {
 		c.unsubscribe(ctx, false, cli, topics...)
 	})
 }
@@ -148,7 +148,7 @@ func (c *RetryClient) removeEstablished(topics ...string) {
 
 // Disconnect from the broker.
 func (c *RetryClient) Disconnect(ctx context.Context) error {
-	err := c.pushTask(ctx, func(cli Client) {
+	err := c.pushTask(ctx, func(ctx context.Context, cli Client) {
 		cli.Disconnect(ctx)
 	})
 	close(c.chTask)
@@ -174,18 +174,19 @@ func (c *RetryClient) SetClient(ctx context.Context, cli Client) {
 		return
 	}
 
-	c.chTask = make(chan func(cli Client))
+	c.chTask = make(chan func(ctx context.Context, cli Client))
 	go func() {
+		ctx := context.Background()
 		for task := range c.chTask {
 			c.mu.Lock()
 			cli := c.cli
 			c.mu.Unlock()
-			task(cli)
+			task(ctx, cli)
 		}
 	}()
 }
 
-func (c *RetryClient) pushTask(ctx context.Context, task func(cli Client)) error {
+func (c *RetryClient) pushTask(ctx context.Context, task func(ctx context.Context, cli Client)) error {
 	c.mu.Lock()
 	chTask := c.chTask
 	c.mu.Unlock()

retryclient_integration_test.go

@@ -54,5 +54,65 @@ func TestIntegration_RetryClient(t *testing.T) {
 			}
 		})
 	}
+}
+
+func TestIntegration_RetryClient_Cancel(t *testing.T) {
+	cliBase, err := Dial(urls["MQTT"], WithTLSConfig(&tls.Config{InsecureSkipVerify: true}))
+	if err != nil {
+		t.Fatalf("Unexpected error: '%v'", err)
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
+	cliRecv, err := Dial(
+		urls["MQTT"],
+		WithTLSConfig(&tls.Config{InsecureSkipVerify: true}),
+	)
+	if err != nil {
+		t.Fatalf("Unexpected error: '%v'", err)
+	}
+	if _, err = cliRecv.Connect(ctx,
+		"RetryClientCancelRecv",
+	); err != nil {
+		t.Fatalf("Unexpected error: '%v'", err)
+	}
+	chRecv := make(chan *Message)
+	cliRecv.Handle(HandlerFunc(func(msg *Message) {
+		chRecv <- msg
+	}))
+	if err := cliRecv.Subscribe(ctx, Subscription{Topic: "testCancel", QoS: QoS2}); err != nil {
+		t.Fatalf("Unexpected error: '%v'", err)
+	}
+
+	var cli RetryClient
+	cli.SetClient(ctx, cliBase)
 
+	if _, err := cli.Connect(ctx, "RetryClientCancel"); err != nil {
+		t.Fatalf("Unexpected error: '%v'", err)
+	}
+
+	ctx2, cancel2 := context.WithCancel(ctx)
+	if err := cli.Publish(ctx2, &Message{
+		Topic:   "testCancel",
+		QoS:     QoS2,
+		Payload: []byte("message"),
+	}); err != nil {
+		t.Fatalf("Unexpected error: '%v'", err)
+	}
+	// Job must be already queued now.
+	cancel2()
+
+	select {
+	case <-chRecv:
+	case <-time.After(time.Second):
+		t.Error("Timeout")
+	}
+
+	if err := cli.Disconnect(ctx); err != nil {
+		t.Fatalf("Unexpected error: '%v'", err)
+	}
+	if err := cliRecv.Disconnect(ctx); err != nil {
+		t.Fatalf("Unexpected error: '%v'", err)
+	}
 }