Merge pull request #191 from rjeczalik/renewer-fix

tokenrenewer: improve fix for a panic #187

Author: rjeczalik

.travis.yml

@@ -1,9 +1,8 @@
 language: go
 sudo: false
 go:
-  - 1.4.3
-  - 1.6.3
-  - 1.7.1
+  - 1.7.4
+  - tip
 install:
   - go get -d -v -t ./...
 script:

client.go

@@ -564,51 +564,55 @@ func (c *Client) OnTokenRenew(handler func(token string)) {
 // callOnConnectHandlers runs the registered connect handlers.
 func (c *Client) callOnConnectHandlers() {
 	c.m.RLock()
+	defer c.m.RUnlock()
+
 	for _, handler := range c.onConnectHandlers {
 		func() {
-			defer recover()
+			defer nopRecover()
 			handler()
 		}()
 	}
-	c.m.RUnlock()
 }
 
 // callOnDisconnectHandlers runs the registered disconnect handlers.
 func (c *Client) callOnDisconnectHandlers() {
 	c.m.RLock()
+	defer c.m.RUnlock()
+
 	for _, handler := range c.onDisconnectHandlers {
 		func() {
-			defer recover()
+			defer nopRecover()
 			handler()
 		}()
 	}
-	c.m.RUnlock()
 }
 
 // callOnTokenExpireHandlers calls registered functions when an error
 // from remote kite is received that token used is expired.
 func (c *Client) callOnTokenExpireHandlers() {
 	c.m.RLock()
+	defer c.m.RUnlock()
+
 	for _, handler := range c.onTokenExpireHandlers {
 		func() {
-			defer recover()
+			defer nopRecover()
 			handler()
 		}()
 	}
-	c.m.RUnlock()
 }
 
 // callOnTokenRenewHandlers calls all registered functions when
 // we successfully obtain new token from kontrol.
 func (c *Client) callOnTokenRenewHandlers(token string) {
 	c.m.RLock()
+	defer c.m.RUnlock()
+
 	for _, handler := range c.onTokenRenewHandlers {
 		func() {
-			defer recover()
+			defer nopRecover()
 			handler(token)
 		}()
 	}
-	c.m.RUnlock()
 }
 
 func (c *Client) wrapMethodArgs(args []interface{}, responseCallback dnode.Function) []interface{} {

handlers.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"net/http"
 	"net/url"
+	"os"
 	"os/exec"
 	"runtime"
 
@@ -66,7 +67,7 @@ func handlePrompt(r *Request) (interface{}, error) {
 // handleGetPass reads a line of input from a terminal without local echo.
 func handleGetPass(r *Request) (interface{}, error) {
 	fmt.Print(r.Args.One().MustString())
-	data, err := terminal.ReadPassword(0) // stdin
+	data, err := terminal.ReadPassword(int(os.Stdin.Fd())) // stdin
 	fmt.Println()
 	if err != nil {
 		return nil, err

kite.go

@@ -324,7 +324,10 @@ func (k *Kite) callOnConnectHandlers(c *Client) {
 	defer k.handlersMu.RUnlock()
 
 	for _, handler := range k.onConnectHandlers {
-		handler(c)
+		func() {
+			defer nopRecover()
+			handler(c)
+		}()
 	}
 }
 
@@ -333,7 +336,10 @@ func (k *Kite) callOnFirstRequestHandlers(c *Client) {
 	defer k.handlersMu.RUnlock()
 
 	for _, handler := range k.onFirstRequestHandlers {
-		handler(c)
+		func() {
+			defer nopRecover()
+			handler(c)
+		}()
 	}
 }
 
@@ -342,7 +348,10 @@ func (k *Kite) callOnDisconnectHandlers(c *Client) {
 	defer k.handlersMu.RUnlock()
 
 	for _, handler := range k.onDisconnectHandlers {
-		handler(c)
+		func() {
+			defer nopRecover()
+			handler(c)
+		}()
 	}
 }
 
@@ -351,7 +360,10 @@ func (k *Kite) callOnRegisterHandlers(r *protocol.RegisterResult) {
 	defer k.handlersMu.RUnlock()
 
 	for _, handler := range k.onRegisterHandlers {
-		handler(r)
+		func() {
+			defer nopRecover()
+			handler(r)
+		}()
 	}
 }
 
@@ -521,3 +533,5 @@ func (fn closerFunc) Close() error {
 
 	return nil
 }
+
+func nopRecover() { recover() }

kite_test.go

@@ -13,6 +13,7 @@ import (
 
 	"github.com/koding/kite/config"
 	"github.com/koding/kite/dnode"
+	"github.com/koding/kite/protocol"
 	"github.com/koding/kite/sockjsclient"
 	_ "github.com/koding/kite/testutil"
 
@@ -23,6 +24,14 @@ func init() {
 	rand.Seed(time.Now().Unix() + int64(os.Getpid()))
 }
 
+func panicHandler(*Client) {
+	panic("this panic should be ignored")
+}
+
+func panicRegisterHandler(*protocol.RegisterResult) {
+	panic("this panic should be ignored")
+}
+
 func TestMultiple(t *testing.T) {
 	testDuration := time.Second * 10
 
@@ -51,6 +60,11 @@ func TestMultiple(t *testing.T) {
 		m.Config.Transport = transport
 		m.Config.Port = port + i
 
+		m.OnConnect(panicHandler)
+		m.OnRegister(panicRegisterHandler)
+		m.OnDisconnect(panicHandler)
+		m.OnFirstRequest(panicHandler)
+
 		m.HandleFunc("square", Square)
 		go m.Run()
 		<-m.ServerReadyNotify()

tokenrenewer.go

@@ -2,7 +2,8 @@ package kite
 
 import (
 	"fmt"
-	"sync/atomic"
+	"strings"
+	"sync"
 	"time"
 
 	jwt "github.com/dgrijalva/jwt-go"
@@ -22,15 +23,16 @@ type TokenRenewer struct {
 	validUntil       time.Time
 	signalRenewToken chan struct{}
 	disconnect       chan struct{}
-	active           uint32
+	once             sync.Once // for c.installHandlers
+	renewLoopWG      sync.WaitGroup
 }
 
 func NewTokenRenewer(r *Client, k *Kite) (*TokenRenewer, error) {
 	t := &TokenRenewer{
 		client:           r,
 		localKite:        k,
-		signalRenewToken: make(chan struct{}, 1),
-		disconnect:       make(chan struct{}, 1),
+		signalRenewToken: make(chan struct{}),
+		disconnect:       make(chan struct{}),
 	}
 	return t, t.parse(r.Auth.Key)
 }
@@ -60,31 +62,35 @@ func (t *TokenRenewer) parse(tokenString string) error {
 
 // RenewWhenExpires renews the token before it expires.
 func (t *TokenRenewer) RenewWhenExpires() {
-	if atomic.CompareAndSwapUint32(&t.active, 0, 1) {
-		t.client.OnConnect(t.startRenewLoop)
-		t.client.OnTokenExpire(t.sendRenewTokenSignal)
-		t.client.OnDisconnect(t.sendDisconnectSignal)
-	}
+	t.once.Do(t.installHandlers)
+}
+
+func (t *TokenRenewer) installHandlers() {
+	t.client.OnConnect(t.startRenewLoop)
+	t.client.OnTokenExpire(t.sendRenewTokenSignal)
+	t.client.OnDisconnect(t.sendDisconnectSignal)
 }
 
 func (t *TokenRenewer) renewLoop() {
+	t.renewLoopWG.Add(1)
+	defer t.renewLoopWG.Done()
+
 	// renews token before it expires (sends the first signal to the goroutine below)
 	go time.AfterFunc(t.renewDuration(), t.sendRenewTokenSignal)
 
 	// renew token on signal util remote kite disconnects.
 	for {
 		select {
 		case <-t.signalRenewToken:
-			switch err := t.renewToken(); err {
-			case nil:
+			switch err := t.renewToken(); {
+			case err == nil:
 				go time.AfterFunc(t.renewDuration(), t.sendRenewTokenSignal)
-			case ErrNoKitesAvailable:
+			case err == ErrNoKitesAvailable || strings.Contains(err.Error(), "no kites found"):
 				// If kite went down we're not going to renew the token,
 				// as we need to dial either way.
 				//
 				// This case handles a situation, when kite missed
 				// disconnect signal (observed to happen with XHR transport).
-				return
 			default:
 				t.localKite.Log.Error("token renewer: %s Cannot renew token for Kite: %s I will retry in %d seconds...",
 					err, t.client.ID, retryInterval/time.Second)
@@ -107,6 +113,14 @@ func (t *TokenRenewer) renewDuration() time.Duration {
 }
 
 func (t *TokenRenewer) startRenewLoop() {
+	// In case when t.client missed a disconnect signal (e.g. due to timeout observed
+	// by the remote end), previous renewLoop will be still running.
+	t.sendDisconnectSignal()
+
+	// if we don't wait to observe previous renewLoop goroutine handle the disconnect
+	// signal, we'd have a race resulting in new renewLoop goroutine handling it.
+	t.renewLoopWG.Wait()
+
 	go t.renewLoop()
 }