Merge 077f583 into 6357932

ndt7.go

@@ -180,26 +180,11 @@ func (c *Client) doConnect(ctx context.Context, serviceURL string) (*websocket.C
 	return conn, err
 }
 
-func (c *Client) getURLforPath(ctx context.Context, p string) (string, error) {
-	// Either the server or service url fields override the Locate API.
-	// First check for the server.
-	if c.Server != "" && (p == params.DownloadURLPath || p == params.UploadURLPath) {
-		return (&url.URL{
-			Scheme: c.Scheme,
-			Host:   c.Server,
-			Path:   p,
-		}).String(), nil
-	}
-	// Second, check for the service url.
-	if c.ServiceURL != nil && (c.ServiceURL.Path == params.DownloadURLPath || c.ServiceURL.Path == params.UploadURLPath) {
-		// Override scheme to match the provided service url.
-		c.Scheme = c.ServiceURL.Scheme
-		return c.ServiceURL.String(), nil
-	} else if c.ServiceURL != nil {
-		return "", ErrServiceUnsupported
-	}
-
-	// Third, neither of the above conditions worked, so use the Locate API.
+// nextURLFromLocate returns the next URL to try from the Locate API.
+// If it's the first time we're calling this function, it contacts the Locate
+// API. Subsequently, it returns the next URL from the cache.
+// If there are no more URLs to try, it returns an error.
+func (c *Client) nextURLFromLocate(ctx context.Context, p string) (string, error) {
 	if len(c.targets) == 0 {
 		targets, err := c.Locate.Nearest(ctx, "ndt/ndt7")
 		if err != nil {
@@ -217,12 +202,9 @@ func (c *Client) getURLforPath(ctx context.Context, p string) (string, error) {
 	return "", ErrNoTargets
 }
 
-// start is the function for starting a test.
-func (c *Client) start(ctx context.Context, f testFn, p string) (<-chan spec.Measurement, error) {
-	s, err := c.getURLforPath(ctx, p)
-	if err != nil {
-		return nil, err
-	}
+// tryConnect tries to establish a websocket connection. If successful, returns
+// a channel where measurements are written.
+func (c *Client) tryConnect(ctx context.Context, f testFn, s string) (<-chan spec.Measurement, error) {
 	u, err := url.Parse(s)
 	if err != nil {
 		return nil, err
@@ -237,6 +219,47 @@ func (c *Client) start(ctx context.Context, f testFn, p string) (<-chan spec.Mea
 	return ch, nil
 }
 
+// start is the function for starting a test.
+func (c *Client) start(ctx context.Context, f testFn, p string) (<-chan spec.Measurement, error) {
+	var customURL *url.URL
+	// Either the server or service url fields override the Locate API.
+	// First check for the server.
+	if c.Server != "" && (p == params.DownloadURLPath || p == params.UploadURLPath) {
+		customURL = &url.URL{
+			Scheme: c.Scheme,
+			Host:   c.Server,
+			Path:   p,
+		}
+	}
+	// Second, check for the service url.
+	if c.ServiceURL != nil && (c.ServiceURL.Path == params.DownloadURLPath || c.ServiceURL.Path == params.UploadURLPath) {
+		// Override scheme to match the provided service url.
+		c.Scheme = c.ServiceURL.Scheme
+		customURL = c.ServiceURL
+	} else if c.ServiceURL != nil {
+		return nil, ErrServiceUnsupported
+	}
+
+	// If a custom URL was provided, use it.
+	if customURL != nil {
+		return c.tryConnect(ctx, f, customURL.String())
+	}
+
+	// If we have no URLs, use the Locate API. In case of failure, try the next
+	// URL until there are no more URLs available.
+	for {
+		s, err := c.nextURLFromLocate(ctx, p)
+		if err != nil {
+			return nil, err
+		}
+		ch, err := c.tryConnect(ctx, f, s)
+		if err != nil {
+			continue
+		}
+		return ch, nil
+	}
+}
+
 func (c *Client) collectData(ctx context.Context, f testFn, conn websocketx.Conn, outch chan<- spec.Measurement) {
 	inch := make(chan spec.Measurement)
 	defer close(outch)

ndt7_test.go

@@ -101,22 +101,19 @@ func TestStartConnectError(t *testing.T) {
 	}
 }
 
-func newLocator(s string) *locatetest.Locator {
-	u, err := url.Parse(s)
-	if err != nil {
-		panic(err)
+// newLocator returns a locate.Client that returns the given server URLs.
+func newLocator(serverURLs []string) *locatetest.Locator {
+	s := []string{}
+	for _, serverURL := range serverURLs {
+		u, err := url.Parse(serverURL)
+		if err != nil {
+			panic(err)
+		}
+		s = append(s, u.Host)
 	}
-	fl := &locatetest.Locator{
-		Servers: []string{u.Host},
+	return &locatetest.Locator{
+		Servers: s,
 	}
-	return fl
-}
-
-func newEmptyLocator() *locatetest.Locator {
-	fl := &locatetest.Locator{
-		Servers: []string{},
-	}
-	return fl
 }
 
 func TestIntegrationDownload(t *testing.T) {
@@ -126,7 +123,7 @@ func TestIntegrationDownload(t *testing.T) {
 	h, fs := ndt7test.NewNDT7Server(t)
 	defer os.RemoveAll(h.DataDir)
 	defer fs.Close()
-	l := locatetest.NewLocateServer(newLocator(fs.URL))
+	l := locatetest.NewLocateServer(newLocator([]string{fs.URL}))
 	client := NewClient(clientName, clientVersion)
 	client.Scheme = "ws"
 	u, err := url.Parse(l.URL + "/v2/nearest")
@@ -160,7 +157,7 @@ func TestIntegrationUpload(t *testing.T) {
 	h, fs := ndt7test.NewNDT7Server(t)
 	defer os.RemoveAll(h.DataDir)
 	defer fs.Close()
-	l := locatetest.NewLocateServer(newLocator(fs.URL))
+	l := locatetest.NewLocateServer(newLocator([]string{fs.URL}))
 	client := NewClient(clientName, clientVersion)
 	client.Scheme = "ws"
 	u, err := url.Parse(l.URL + "/v2/nearest")
@@ -241,7 +238,8 @@ func TestDownloadNoTargets(t *testing.T) {
 	h, fs := ndt7test.NewNDT7Server(t)
 	defer os.RemoveAll(h.DataDir)
 	defer fs.Close()
-	l := locatetest.NewLocateServer(newLocator(fs.URL))
+	// The first URL is intentionally invalid to test the retry loop.
+	l := locatetest.NewLocateServer(newLocator([]string{"https://invalid", fs.URL}))
 	client := NewClient(clientName, clientVersion)
 	client.Scheme = "ws"
 	u, err := url.Parse(l.URL + "/v2/nearest")
@@ -250,7 +248,7 @@ func TestDownloadNoTargets(t *testing.T) {
 	loc.BaseURL = u
 	client.Locate = loc
 
-	// First attempt should succeed.
+	// This should succeed by using the second URL since the first one fails.
 	ch, err := client.StartDownload(context.Background())
 	testingx.Must(t, err, "failed to download first attempt")
 	tot := 0
@@ -260,7 +258,8 @@ func TestDownloadNoTargets(t *testing.T) {
 	if tot <= 0 {
 		t.Fatal("Expected at least a measurement")
 	}
-	// Second attempt should return ErrNoTargets.
+	// Second attempt should return ErrNoTargets since all the available
+	// servers have been tried.
 	_, err = client.StartDownload(context.Background())
 	if err != ErrNoTargets {
 		t.Fatalf("Expected no target error: %v", err)