kite: cancel Context when client disconnects

Author: rjeczalik

client.go

@@ -1,6 +1,7 @@
 package kite
 
 import (
+	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -126,6 +127,11 @@ type Client struct {
 	session sockjs.Session
 	send    chan *message
 
+	// ctx and cancel keeps track of session lifetime
+	ctxMu  sync.Mutex
+	ctx    context.Context
+	cancel func()
+
 	// muReconnect protects Reconnect
 	muReconnect sync.Mutex
 
@@ -206,8 +212,13 @@ func (k *Kite) NewClient(remoteURL string) *Client {
 		Concurrent:         true,
 		send:               make(chan *message),
 		interrupt:          make(chan error, 1),
+		ctx:                context.TODO(),
+		cancel:             func() {},
 	}
 
+	c.OnConnect(c.setContext)
+	c.OnDisconnect(c.closeContext)
+
 	k.OnRegister(c.updateAuth)
 
 	return c
@@ -262,6 +273,24 @@ func (c *Client) updateAuth(reg *protocol.RegisterResult) {
 	}
 }
 
+func (c *Client) setContext() {
+	c.ctxMu.Lock()
+	c.ctx, c.cancel = context.WithCancel(context.Background())
+	c.ctxMu.Unlock()
+}
+
+func (c *Client) closeContext() {
+	c.ctxMu.Lock()
+	c.cancel()
+	c.ctxMu.Unlock()
+}
+
+func (c *Client) context() context.Context {
+	c.ctxMu.Lock()
+	defer c.ctxMu.Unlock()
+	return c.ctx
+}
+
 func (c *Client) authCopy() *Auth {
 	c.authMu.Lock()
 	defer c.authMu.Unlock()

kite.go

@@ -283,6 +283,7 @@ func (k *Kite) sockjsHandler(session sockjs.Session) {
 	go c.sendHub()
 
 	k.callOnConnectHandlers(c)
+	c.callOnConnectHandlers()
 
 	// Run after methods are registered and delegate is set
 	c.readLoop()

kite_test.go

@@ -2,10 +2,12 @@ package kite
 
 import (
 	"errors"
+	"flag"
 	"fmt"
 	"math"
 	"math/rand"
 	"os"
+	"reflect"
 	"strconv"
 	"sync"
 	"testing"
@@ -20,6 +22,8 @@ import (
 	"github.com/igm/sockjs-go/sockjs"
 )
 
+var timeout = flag.Duration("telltime", 4*time.Second, "Timeout for kite calls.")
+
 func init() {
 	rand.Seed(time.Now().Unix() + int64(os.Getpid()))
 }
@@ -32,6 +36,76 @@ func panicRegisterHandler(*protocol.RegisterResult) {
 	panic("this panic should be ignored")
 }
 
+func transportFromEnv() config.Transport {
+	env := os.Getenv("KITE_TRANSPORT")
+	tr, ok := config.Transports[env]
+	if env != "" && !ok {
+		panic(fmt.Errorf("transport %q doesn't exists", env))
+	}
+	return tr
+}
+
+func TestContext(t *testing.T) {
+	flag.Parse()
+
+	ch := make(chan int, 4) // checkpoints, to ensure flor of control
+
+	k := New("server", "0.0.1")
+	k.Config.DisableAuthentication = true
+	k.Config.Port = 3333
+	k.Config.Transport = transportFromEnv()
+	k.HandleFunc("longrunning", func(r *Request) (interface{}, error) {
+		ch <- 2
+
+		go func() {
+			<-r.Context.Done()
+			ch <- 4
+			close(ch)
+		}()
+		return nil, nil
+	})
+	go k.Run()
+	<-k.ServerReadyNotify()
+	defer k.Close()
+
+	c := New("client", "0.0.1").NewClient("http://127.0.0.1:3333/kite")
+	if err := c.Dial(); err != nil {
+		t.Fatalf("Dial()=%s", err)
+	}
+
+	ch <- 1
+
+	if _, err := c.TellWithTimeout("longrunning", *timeout); err != nil {
+		t.Fatalf("TellWithTimeout()=%s", err)
+	}
+
+	ch <- 3
+
+	c.Close()
+
+	var got []int
+	timeout := time.After(2 * time.Second)
+
+collect:
+	for {
+		select {
+		case i, ok := <-ch:
+			if !ok {
+				break collect
+			}
+			got = append(got, i)
+		case <-timeout:
+			t.Fatal("timed out collecting checkpoints")
+		}
+	}
+
+	want := []int{1, 2, 3, 4}
+
+	if !reflect.DeepEqual(got, want) {
+		t.Fatalf("got %v, want %v")
+	}
+}
+
 func TestMultiple(t *testing.T) {
 	testDuration := time.Second * 10
 
@@ -44,15 +118,7 @@ func TestMultiple(t *testing.T) {
 	// ports are starting from 6000 up to 6000 + kiteNumber
 	port := 6000
 
-	var transport config.Transport
-	if transportName := os.Getenv("KITE_TRANSPORT"); transportName != "" {
-		tr, ok := config.Transports[transportName]
-		if !ok {
-			t.Fatalf("transport '%s' doesn't exists", transportName)
-		}
-
-		transport = tr
-	}
+	transport := transportFromEnv()
 
 	for i := 0; i < kiteNumber; i++ {
 		m := New("mathworker"+strconv.Itoa(i), "0.1."+strconv.Itoa(i))

request.go

@@ -154,7 +154,7 @@ func (c *Client) newRequest(method string, args *dnode.Partial) (*Request, func(
 		LocalKite: c.LocalKite,
 		Client:    c,
 		Auth:      options.Auth,
-		Context:   context.TODO(),
+		Context:   c.context(),
 	}
 
 	// Call response callback function, send back our response