Merge lp:~pedronis/ubuntu-push/try-hosts into lp:ubuntu-push

[Samuele Pedroni]

optionally retrieves a list of hosts to try to connect to from a url,

still need changes in the client config and the interface between client and session: for a second branch, this one is already biggish

found that session.Close() has a race about Connection, will fix also in a later branch

[lp:~pedronis/ubuntu-push/try-hosts]

98. By Samuele Pedroni on 2014-03-28

clarity

99. By Samuele Pedroni on 2014-03-28

ordering

100. By Samuele Pedroni on 2014-03-28

formatting

Unmerged revisions

=== modified file 'client/session/session.go'
--- client/session/session.go	2014-03-26 16:26:36 +0000
+++ client/session/session.go	2014-03-27 21:20:49 +0000
@@ -23,14 +23,17 @@
 	"crypto/x509"
 	"errors"
 	"fmt"
+	"math/rand"
+	"net"
+	"strings"
+	"sync/atomic"
+	"time"
+
+	"launchpad.net/ubuntu-push/client/gethosts"
 	"launchpad.net/ubuntu-push/client/session/levelmap"
 	"launchpad.net/ubuntu-push/logger"
 	"launchpad.net/ubuntu-push/protocol"
 	"launchpad.net/ubuntu-push/util"
-	"math/rand"
-	"net"
-	"sync/atomic"
-	"time"
 )
 
 var wireVersionBytes = []byte{protocol.ProtocolWireVersion}
@@ -45,6 +48,15 @@
 	protocol.NotificationsMsg
 }
 
+// parseServerAddrSpec recognizes whether spec is a HTTP URL to get
+// hosts from or a |-separated list of host:port pairs.
+func parseServerAddrSpec(spec string) (hostsEndpoint string, fallbackHosts []string) {
+	if strings.HasPrefix(spec, "http") {
+		return spec, nil
+	}
+	return "", strings.Split(spec, "|")
+}
+
 // ClientSessionState is a way to broadly track the progress of the session
 type ClientSessionState uint32
 
@@ -56,14 +68,27 @@
 	Running
 )
 
+type hostGetter interface {
+	Get() ([]string, error)
+}
+
 // ClienSession holds a client<->server session and its configuration.
 type ClientSession struct {
 	// configuration
-	DeviceId        string
-	ServerAddr      string
-	ExchangeTimeout time.Duration
-	Levels          levelmap.LevelMap
-	Protocolator    func(net.Conn) protocol.Protocol
+	DeviceId           string
+	getHost            hostGetter
+	fallbackHosts      []string
+	ExchangeTimeout    time.Duration
+	HostsCachingExpiry time.Duration
+	Levels             levelmap.LevelMap
+	Protocolator       func(net.Conn) protocol.Protocol
+	// hosts
+	timeSince              func(time.Time) time.Duration // hook for testing
+	deliveryHostsTimestamp time.Time
+	deliveryHosts          []string
+	lastAttemptTimestamp   time.Time
+	leftToTry              int
+	tryHost                int
 	// connection
 	Connection   net.Conn
 	Log          logger.Logger
@@ -77,7 +102,7 @@
 	MsgCh  chan *Notification
 }
 
-func NewSession(serverAddr string, pem []byte, exchangeTimeout time.Duration,
+func NewSession(serverAddrSpec string, pem []byte, exchangeTimeout time.Duration,
 	deviceId string, levelmapFactory func() (levelmap.LevelMap, error),
 	log logger.Logger) (*ClientSession, error) {
 	state := uint32(Disconnected)
@@ -85,15 +110,23 @@
 	if err != nil {
 		return nil, err
 	}
+	var getHost hostGetter
+	hostsEndpoint, fallbackHosts := parseServerAddrSpec(serverAddrSpec)
+	if hostsEndpoint != "" {
+		getHost = gethosts.New(deviceId, hostsEndpoint, exchangeTimeout)
+	}
 	sess := &ClientSession{
-		ExchangeTimeout: exchangeTimeout,
-		ServerAddr:      serverAddr,
-		DeviceId:        deviceId,
-		Log:             log,
-		Protocolator:    protocol.NewProtocol0,
-		Levels:          levels,
-		TLS:             &tls.Config{InsecureSkipVerify: true}, // XXX
-		stateP:          &state,
+		ExchangeTimeout:    exchangeTimeout,
+		HostsCachingExpiry: 12 * time.Hour,
+		getHost:            getHost,
+		fallbackHosts:      fallbackHosts,
+		DeviceId:           deviceId,
+		Log:                log,
+		Protocolator:       protocol.NewProtocol0,
+		Levels:             levels,
+		TLS:                &tls.Config{InsecureSkipVerify: true}, // XXX
+		stateP:             &state,
+		timeSince:          time.Since,
 	}
 	if pem != nil {
 		cp := x509.NewCertPool()
@@ -114,13 +147,72 @@
 	atomic.StoreUint32(sess.stateP, uint32(state))
 }
 
+// getHosts sets deliverHosts possibly querying a remote endpoint
+func (sess *ClientSession) getHosts() error {
+	if sess.getHost != nil {
+		if sess.timeSince(sess.deliveryHostsTimestamp) < sess.HostsCachingExpiry {
+			return nil
+		}
+		hosts, err := sess.getHost.Get()
+		if err != nil {
+			sess.Log.Errorf("getHosts: %v", err)
+			sess.setState(Error)
+			return err
+		}
+		sess.deliveryHostsTimestamp = time.Now()
+		sess.deliveryHosts = hosts
+	} else {
+		sess.deliveryHosts = sess.fallbackHosts
+	}
+	return nil
+}
+
+// startConnectionAttempt/nextHostToTry help connect iterating over candidate hosts
+
+func (sess *ClientSession) startConnectionAttempt() {
+	if sess.timeSince(sess.lastAttemptTimestamp) > 10*sess.ExchangeTimeout {
+		sess.tryHost = 0
+	}
+	sess.leftToTry = len(sess.deliveryHosts)
+	sess.lastAttemptTimestamp = time.Now()
+}
+
+func (sess *ClientSession) nextHostToTry() string {
+	if sess.leftToTry == 0 {
+		return ""
+	}
+	res := sess.deliveryHosts[sess.tryHost]
+	sess.tryHost = (sess.tryHost + 1) % len(sess.deliveryHosts)
+	sess.leftToTry--
+	return res
+}
+
+// we reached the Started state, we can retry with the same host if we
+// have to retry again
+func (sess *ClientSession) started() {
+	sess.tryHost--
+	if sess.tryHost == -1 {
+		sess.tryHost = len(sess.deliveryHosts) - 1
+	}
+	sess.setState(Started)
+}
+
 // connect to a server using the configuration in the ClientSession
 // and set up the connection.
 func (sess *ClientSession) connect() error {
-	conn, err := net.DialTimeout("tcp", sess.ServerAddr, sess.ExchangeTimeout)
-	if err != nil {
-		sess.setState(Error)
-		return fmt.Errorf("connect: %s", err)
+	sess.startConnectionAttempt()
+	var err error
+	var conn net.Conn
+	for {
+		host := sess.nextHostToTry()
+		if host == "" {
+			sess.setState(Error)
+			return fmt.Errorf("connect: %s", err)
+		}
+		conn, err = net.DialTimeout("tcp", host, sess.ExchangeTimeout)
+		if err == nil {
+			break
+		}
 	}
 	sess.Connection = tls.Client(conn, sess.TLS)
 	sess.setState(Connected)
@@ -279,15 +371,19 @@
 	sess.proto = proto
 	sess.pingInterval = pingInterval
 	sess.Log.Debugf("Connected %v.", conn.LocalAddr())
-	sess.setState(Started)
+	sess.started() // deals with choosing which host to retry with as well
 	return nil
 }
 
 // run calls connect, and if it works it calls start, and if it works
 // it runs loop in a goroutine, and ships its return value over ErrCh.
-func (sess *ClientSession) run(closer func(), connecter, starter, looper func() error) error {
+func (sess *ClientSession) run(closer func(), hostGetter, connecter, starter, looper func() error) error {
 	closer()
-	err := connecter()
+	err := hostGetter()
+	if err != nil {
+		return err
+	}
+	err = connecter()
 	if err == nil {
 		err = starter()
 		if err == nil {
@@ -317,7 +413,7 @@
 		// keep on trying.
 		panic("can't Dial() without a protocol constructor.")
 	}
-	return sess.run(sess.doClose, sess.connect, sess.start, sess.loop)
+	return sess.run(sess.doClose, sess.getHosts, sess.connect, sess.start, sess.loop)
 }
 
 func init() {

=== modified file 'client/session/session_test.go'
--- client/session/session_test.go	2014-03-27 13:26:10 +0000
+++ client/session/session_test.go	2014-03-27 21:20:49 +0000
@@ -23,16 +23,21 @@
 	"fmt"
 	"io"
 	"io/ioutil"
+	"net"
+	"net/http"
+	"net/http/httptest"
+	"reflect"
+	"testing"
+	"time"
+
 	. "launchpad.net/gocheck"
+
 	"launchpad.net/ubuntu-push/client/session/levelmap"
+	//"launchpad.net/ubuntu-push/client/gethosts"
 	"launchpad.net/ubuntu-push/logger"
 	"launchpad.net/ubuntu-push/protocol"
 	helpers "launchpad.net/ubuntu-push/testing"
 	"launchpad.net/ubuntu-push/testing/condition"
-	"net"
-	"reflect"
-	"testing"
-	"time"
 )
 
 func TestSession(t *testing.T) { TestingT(t) }
@@ -181,18 +186,43 @@
 }
 
 /****************************************************************
+  parseServerAddrSpec() tests
+****************************************************************/
+
+func (cs *clientSessionSuite) TestParseServerAddrSpec(c *C) {
+	hEp, fallbackHosts := parseServerAddrSpec("http://foo/hosts")
+	c.Check(hEp, Equals, "http://foo/hosts")
+	c.Check(fallbackHosts, IsNil)
+
+	hEp, fallbackHosts = parseServerAddrSpec("foo:443")
+	c.Check(hEp, Equals, "")
+	c.Check(fallbackHosts, DeepEquals, []string{"foo:443"})
+
+	hEp, fallbackHosts = parseServerAddrSpec("foo:443|bar:443")
+	c.Check(hEp, Equals, "")
+	c.Check(fallbackHosts, DeepEquals, []string{"foo:443", "bar:443"})
+}
+
+/****************************************************************
   NewSession() tests
 ****************************************************************/
 
 func (cs *clientSessionSuite) TestNewSessionPlainWorks(c *C) {
-	sess, err := NewSession("", nil, 0, "", cs.lvls, cs.log)
+	sess, err := NewSession("foo:443", nil, 0, "", cs.lvls, cs.log)
 	c.Check(sess, NotNil)
 	c.Check(err, IsNil)
+	c.Check(sess.fallbackHosts, DeepEquals, []string{"foo:443"})
 	// but no root CAs set
 	c.Check(sess.TLS.RootCAs, IsNil)
 	c.Check(sess.State(), Equals, Disconnected)
 }
 
+func (cs *clientSessionSuite) TestNewSessionHostEndpointWorks(c *C) {
+	sess, err := NewSession("http://foo/hosts", pem, 0, "wah", cs.lvls, cs.log)
+	c.Assert(err, IsNil)
+	c.Check(sess.getHost, NotNil)
+}
+
 var certfile string = helpers.SourceRelative("../../server/acceptance/config/testing.cert")
 var pem, _ = ioutil.ReadFile(certfile)
 
@@ -218,12 +248,141 @@
 }
 
 /****************************************************************
+  getHosts() tests
+****************************************************************/
+
+func (cs *clientSessionSuite) TestGetHostsFallback(c *C) {
+	fallback := []string{"foo:443", "bar:443"}
+	sess := &ClientSession{fallbackHosts: fallback}
+	err := sess.getHosts()
+	c.Assert(err, IsNil)
+	c.Check(sess.deliveryHosts, DeepEquals, fallback)
+}
+
+type testHostGetter struct {
+	hosts []string
+	err   error
+}
+
+func (thg *testHostGetter) Get() ([]string, error) {
+	return thg.hosts, thg.err
+}
+
+func (cs *clientSessionSuite) TestGetHostsRemote(c *C) {
+	hostGetter := &testHostGetter{[]string{"foo:443", "bar:443"}, nil}
+	sess := &ClientSession{getHost: hostGetter, timeSince: time.Since}
+	err := sess.getHosts()
+	c.Assert(err, IsNil)
+	c.Check(sess.deliveryHosts, DeepEquals, []string{"foo:443", "bar:443"})
+}
+
+func (cs *clientSessionSuite) TestGetHostsRemoteError(c *C) {
+	sess, err := NewSession("", nil, 0, "", cs.lvls, cs.log)
+	c.Assert(err, IsNil)
+	hostsErr := errors.New("failed")
+	hostGetter := &testHostGetter{nil, hostsErr}
+	sess.getHost = hostGetter
+	err = sess.getHosts()
+	c.Assert(err, Equals, hostsErr)
+	c.Check(sess.deliveryHosts, IsNil)
+	c.Check(sess.State(), Equals, Error)
+}
+
+func (cs *clientSessionSuite) TestGetHostsRemoteCaching(c *C) {
+	hostGetter := &testHostGetter{[]string{"foo:443", "bar:443"}, nil}
+	sess := &ClientSession{
+		getHost:            hostGetter,
+		HostsCachingExpiry: 2 * time.Hour,
+		timeSince:          time.Since,
+	}
+	err := sess.getHosts()
+	c.Assert(err, IsNil)
+	hostGetter.hosts = []string{"baz:443"}
+	// cached
+	err = sess.getHosts()
+	c.Assert(err, IsNil)
+	c.Check(sess.deliveryHosts, DeepEquals, []string{"foo:443", "bar:443"})
+	// expired
+	sess.timeSince = func(ts time.Time) time.Duration {
+		return 3 * time.Hour
+	}
+	err = sess.getHosts()
+	c.Assert(err, IsNil)
+	c.Check(sess.deliveryHosts, DeepEquals, []string{"baz:443"})
+}
+
+/****************************************************************
+  startConnectionAttempt()/nextHostToTry()/started tests
+****************************************************************/
+
+func (cs *clientSessionSuite) TestStartConnectionAttempt(c *C) {
+	since := time.Since(time.Time{})
+	sess := &ClientSession{
+		ExchangeTimeout: 10 * time.Second,
+		timeSince: func(ts time.Time) time.Duration {
+			return since
+		},
+		deliveryHosts: []string{"foo:443", "bar:443"},
+	}
+	// start from first host
+	sess.startConnectionAttempt()
+	c.Check(sess.lastAttemptTimestamp, Not(Equals), 0)
+	c.Check(sess.tryHost, Equals, 0)
+	c.Check(sess.leftToTry, Equals, 2)
+	since = 1 * time.Second
+	sess.tryHost = 1
+	// just continue
+	sess.startConnectionAttempt()
+	c.Check(sess.tryHost, Equals, 1)
+	sess.tryHost = 2
+}
+
+func (cs *clientSessionSuite) TestNextHostToTry(c *C) {
+	sess := &ClientSession{
+		deliveryHosts: []string{"foo:443", "bar:443", "baz:443"},
+		tryHost:       0,
+		leftToTry:     3,
+	}
+	c.Check(sess.nextHostToTry(), Equals, "foo:443")
+	c.Check(sess.nextHostToTry(), Equals, "bar:443")
+	c.Check(sess.nextHostToTry(), Equals, "baz:443")
+	c.Check(sess.nextHostToTry(), Equals, "")
+	c.Check(sess.nextHostToTry(), Equals, "")
+	c.Check(sess.tryHost, Equals, 0)
+
+	sess.leftToTry = 3
+	sess.tryHost = 1
+	c.Check(sess.nextHostToTry(), Equals, "bar:443")
+	c.Check(sess.nextHostToTry(), Equals, "baz:443")
+	c.Check(sess.nextHostToTry(), Equals, "foo:443")
+	c.Check(sess.nextHostToTry(), Equals, "")
+	c.Check(sess.nextHostToTry(), Equals, "")
+	c.Check(sess.tryHost, Equals, 1)
+}
+
+func (cs *clientSessionSuite) TestStarted(c *C) {
+	sess, err := NewSession("", nil, 0, "", cs.lvls, cs.log)
+	c.Assert(err, IsNil)
+
+	sess.deliveryHosts = []string{"foo:443", "bar:443", "baz:443"}
+	sess.tryHost = 1
+
+	sess.started()
+	c.Check(sess.tryHost, Equals, 0)
+	c.Check(sess.State(), Equals, Started)
+
+	sess.started()
+	c.Check(sess.tryHost, Equals, 2)
+}
+
+/****************************************************************
   connect() tests
 ****************************************************************/
 
 func (cs *clientSessionSuite) TestConnectFailsWithNoAddress(c *C) {
 	sess, err := NewSession("", nil, 0, "wah", cs.lvls, cs.log)
 	c.Assert(err, IsNil)
+	sess.deliveryHosts = []string{"nowhere"}
 	err = sess.connect()
 	c.Check(err, ErrorMatches, ".*connect.*address.*")
 	c.Check(sess.State(), Equals, Error)
@@ -233,12 +392,27 @@
 	srv, err := net.Listen("tcp", "localhost:0")
 	c.Assert(err, IsNil)
 	defer srv.Close()
-	sess, err := NewSession(srv.Addr().String(), nil, 0, "wah", cs.lvls, cs.log)
-	c.Assert(err, IsNil)
-	err = sess.connect()
-	c.Check(err, IsNil)
-	c.Check(sess.Connection, NotNil)
-	c.Check(sess.State(), Equals, Connected)
+	sess, err := NewSession("", nil, 0, "wah", cs.lvls, cs.log)
+	c.Assert(err, IsNil)
+	sess.deliveryHosts = []string{srv.Addr().String()}
+	err = sess.connect()
+	c.Check(err, IsNil)
+	c.Check(sess.Connection, NotNil)
+	c.Check(sess.State(), Equals, Connected)
+}
+
+func (cs *clientSessionSuite) TestConnectSecondConnects(c *C) {
+	srv, err := net.Listen("tcp", "localhost:0")
+	c.Assert(err, IsNil)
+	defer srv.Close()
+	sess, err := NewSession("", nil, 0, "wah", cs.lvls, cs.log)
+	c.Assert(err, IsNil)
+	sess.deliveryHosts = []string{"nowhere", srv.Addr().String()}
+	err = sess.connect()
+	c.Check(err, IsNil)
+	c.Check(sess.Connection, NotNil)
+	c.Check(sess.State(), Equals, Connected)
+	c.Check(sess.tryHost, Equals, 0)
 }
 
 func (cs *clientSessionSuite) TestConnectConnectFail(c *C) {
@@ -247,6 +421,7 @@
 	sess, err := NewSession(srv.Addr().String(), nil, 0, "wah", cs.lvls, cs.log)
 	srv.Close()
 	c.Assert(err, IsNil)
+	sess.deliveryHosts = []string{srv.Addr().String()}
 	err = sess.connect()
 	c.Check(err, ErrorMatches, ".*connection refused")
 	c.Check(sess.State(), Equals, Error)
@@ -688,20 +863,34 @@
   run() tests
 ****************************************************************/
 
-func (cs *clientSessionSuite) TestRunBailsIfConnectFails(c *C) {
+func (cs *clientSessionSuite) TestRunBailsIfHostGetterFails(c *C) {
 	sess, err := NewSession("", nil, 0, "wah", cs.lvls, cs.log)
 	c.Assert(err, IsNil)
-	failure := errors.New("TestRunBailsIfConnectFails")
+	failure := errors.New("TestRunBailsIfHostGetterFails")
 	has_closed := false
 	err = sess.run(
 		func() { has_closed = true },
 		func() error { return failure },
 		nil,
+		nil,
 		nil)
 	c.Check(err, Equals, failure)
 	c.Check(has_closed, Equals, true)
 }
 
+func (cs *clientSessionSuite) TestRunBailsIfConnectFails(c *C) {
+	sess, err := NewSession("", nil, 0, "wah", cs.lvls, cs.log)
+	c.Assert(err, IsNil)
+	failure := errors.New("TestRunBailsIfConnectFails")
+	err = sess.run(
+		func() {},
+		func() error { return nil },
+		func() error { return failure },
+		nil,
+		nil)
+	c.Check(err, Equals, failure)
+}
+
 func (cs *clientSessionSuite) TestRunBailsIfStartFails(c *C) {
 	sess, err := NewSession("", nil, 0, "wah", cs.lvls, cs.log)
 	c.Assert(err, IsNil)
@@ -709,6 +898,7 @@
 	err = sess.run(
 		func() {},
 		func() error { return nil },
+		func() error { return nil },
 		func() error { return failure },
 		nil)
 	c.Check(err, Equals, failure)
@@ -727,6 +917,7 @@
 		func() {},
 		func() error { return nil },
 		func() error { return nil },
+		func() error { return nil },
 		func() error { sess.MsgCh <- notf; return <-failureCh })
 	c.Check(err, Equals, nil)
 	// if run doesn't error it sets up the channels
@@ -794,7 +985,20 @@
 	timeout := 100 * time.Millisecond
 	lst, err := tls.Listen("tcp", "localhost:0", tlsCfg)
 	c.Assert(err, IsNil)
-	sess, err := NewSession(lst.Addr().String(), nil, timeout, "wah", cs.lvls, cs.log)
+	// advertise
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		b, err := json.Marshal(map[string]interface{}{
+			"hosts": []string{"nowhere", lst.Addr().String()},
+		})
+		if err != nil {
+			panic(err)
+		}
+		w.Header().Set("Content-Type", "application/json")
+		w.Write(b)
+	}))
+	defer ts.Close()
+
+	sess, err := NewSession(ts.URL, nil, timeout, "wah", cs.lvls, cs.log)
 	c.Assert(err, IsNil)
 	tconn := &testConn{CloseCondition: condition.Fail2Work(10)}
 	sess.Connection = tconn
@@ -823,6 +1027,9 @@
 	c.Assert(err, IsNil)
 	c.Assert(v, Equals, protocol.ProtocolWireVersion)
 
+	// if something goes wrong session would try the first/other host
+	c.Check(sess.tryHost, Equals, 0)
+
 	// 2. "connect" (but on the fake protcol above! woo)
 
 	c.Check(takeNext(downCh), Equals, "deadline 100ms")
@@ -843,6 +1050,9 @@
 	c.Check(takeNext(downCh), Equals, protocol.PingPongMsg{Type: "pong"})
 	upCh <- nil
 
+	// session would retry the same host
+	c.Check(sess.tryHost, Equals, 1)
+
 	// and broadcasts...
 	b := &protocol.BroadcastMsg{
 		Type:     "broadcast",
@@ -870,3 +1080,31 @@
 	upCh <- failure
 	c.Check(<-sess.ErrCh, Equals, failure)
 }
+
+func (cs *clientSessionSuite) TestDialWorksDirect(c *C) {
+	// happy path thoughts
+	cert, err := tls.X509KeyPair(helpers.TestCertPEMBlock, helpers.TestKeyPEMBlock)
+	c.Assert(err, IsNil)
+	tlsCfg := &tls.Config{
+		Certificates:           []tls.Certificate{cert},
+		SessionTicketsDisabled: true,
+	}
+
+	timeout := 100 * time.Millisecond
+	lst, err := tls.Listen("tcp", "localhost:0", tlsCfg)
+	c.Assert(err, IsNil)
+	sess, err := NewSession(lst.Addr().String(), nil, timeout, "wah", cs.lvls, cs.log)
+	c.Assert(err, IsNil)
+	//defer sess.Close() xxx provokes a race, fix in a later branch
+
+	upCh := make(chan interface{}, 5)
+	downCh := make(chan interface{}, 5)
+	proto := &testProtocol{up: upCh, down: downCh}
+	sess.Protocolator = func(net.Conn) protocol.Protocol { return proto }
+
+	go sess.Dial()
+
+	_, err = lst.Accept()
+	c.Assert(err, IsNil)
+	// connect done
+}