Enable polling in the source domains

Enables gitlabClient.Poll in the domain source init.

10 files changed, 87 insertions, 33 deletions

diff --git a/internal/source/domains.go b/internal/source/domains.go
index 77e1aa1d..9bb5f38e 100644
--- a/internal/source/domains.go
+++ b/internal/source/domains.go
@@ -3,6 +3,7 @@ package source
 import (
 	"errors"
 	"regexp"
+	"sync"
 	"time"
 
 	log "github.com/sirupsen/logrus"
@@ -12,6 +13,7 @@ import (
 	"gitlab.com/gitlab-org/gitlab-pages/internal/source/disk"
 	"gitlab.com/gitlab-org/gitlab-pages/internal/source/domains/gitlabsourceconfig"
 	"gitlab.com/gitlab-org/gitlab-pages/internal/source/gitlab"
+	"gitlab.com/gitlab-org/gitlab-pages/internal/source/gitlab/client"
 )
 
 var (
@@ -33,6 +35,7 @@ func init() {
 // currently using two sources during the transition to the new GitLab domains
 // source.
 type Domains struct {
+	mu     *sync.RWMutex
 	gitlab Source
 	disk   *disk.Disk // legacy disk source
 }
@@ -44,19 +47,47 @@ func NewDomains(config Config) (*Domains, error) {
 	// TODO: choose domain source config via config.DomainConfigSource()
 	// https://gitlab.com/gitlab-org/gitlab/-/issues/217912
 
+	domains := &Domains{
+		mu:   &sync.RWMutex{},
+		disk: disk.New(),
+	}
+
 	if len(config.InternalGitLabServerURL()) == 0 || len(config.GitlabAPISecret()) == 0 {
-		return &Domains{disk: disk.New()}, nil
+		return domains, nil
 	}
 
-	gitlab, err := gitlab.New(config)
+	gitlabClient, err := gitlab.New(config)
 	if err != nil {
 		return nil, err
 	}
+	gitlabErr := make(chan error)
+	gitlabClient.Poll(client.DefaultPollingMaxRetries, client.DefaultPollingInterval, gitlabErr)
+
+	domains.enableGitLabSource(gitlabClient)
+
+	go func() {
+		err := <-gitlabErr
+		if err != nil {
+			log.WithError(err).Error("failed to connect to the GitLab API")
+			domains.disableGitLabSource()
+		}
+	}()
+
+	return domains, nil
+}
+
+func (d *Domains) enableGitLabSource(gitlabClient *gitlab.Gitlab) {
+	d.mu.Lock()
+	defer d.mu.Unlock()
+
+	d.gitlab = gitlabClient
+}
+
+func (d *Domains) disableGitLabSource() {
+	d.mu.Lock()
+	defer d.mu.Unlock()
 
-	return &Domains{
-		gitlab: gitlab,
-		disk:   disk.New(),
-	}, nil
+	d.gitlab = nil
 }
 
 // GetDomain retrieves a domain information from a source. We are using two
@@ -85,6 +116,9 @@ func (d *Domains) IsReady() bool {
 }
 
 func (d *Domains) source(domain string) Source {
+	d.mu.RLock()
+	defer d.mu.RUnlock()
+
 	if d.gitlab == nil {
 		return d.disk
 	}
diff --git a/internal/source/domains_test.go b/internal/source/domains_test.go
index 24008b08..da6f3877 100644
--- a/internal/source/domains_test.go
+++ b/internal/source/domains_test.go
@@ -2,6 +2,7 @@ package source
 
 import (
 	"math/rand"
+	"sync"
 	"testing"
 	"time"
 
@@ -67,10 +68,7 @@ func TestGetDomain(t *testing.T) {
 			Once()
 		defer newSource.AssertExpectations(t)
 
-		domains := &Domains{
-			disk:   disk.New(),
-			gitlab: newSource,
-		}
+		domains := newTestDomains(t, newSource)
 
 		domains.GetDomain(testDomain)
 	})
@@ -79,10 +77,7 @@ func TestGetDomain(t *testing.T) {
 		newSource := NewMockSource()
 		defer newSource.AssertExpectations(t)
 
-		domains := &Domains{
-			disk:   disk.New(),
-			gitlab: newSource,
-		}
+		domains := newTestDomains(t, newSource)
 
 		domain, err := domains.GetDomain("domain.test.io")
 
@@ -94,10 +89,7 @@ func TestGetDomain(t *testing.T) {
 		newSource := NewMockSource()
 		defer newSource.AssertExpectations(t)
 
-		domains := &Domains{
-			disk:   disk.New(),
-			gitlab: newSource,
-		}
+		domains := newTestDomains(t, newSource)
 
 		domain, err := domains.GetDomain("pages-broken-poc.gitlab.io")
 
@@ -124,10 +116,7 @@ func TestGetDomain(t *testing.T) {
 			Once()
 		defer newSource.AssertExpectations(t)
 
-		domains := &Domains{
-			disk:   disk.New(),
-			gitlab: newSource,
-		}
+		domains := newTestDomains(t, newSource)
 
 		domains.GetDomain(testDomain)
 	})
@@ -156,8 +145,6 @@ func TestGetDomainWithIncrementalrolloutOfGitLabSource(t *testing.T) {
 	// Generates FNV 2643293380, 2643293380 % 100 = 80
 	domain80 := "test-domain-b.com"
 
-	diskSource := disk.New()
-
 	gitlabSourceConfig.Domains.Rollout.Percentage = 80
 
 	type testDomain struct {
@@ -203,10 +190,7 @@ func TestGetDomainWithIncrementalrolloutOfGitLabSource(t *testing.T) {
 			}
 			defer gitlabSource.AssertExpectations(t)
 
-			domains := &Domains{
-				disk:   diskSource,
-				gitlab: gitlabSource,
-			}
+			domains := newTestDomains(t, gitlabSource)
 
 			gitlabSourceConfig.Domains.Rollout.Stickiness = tc.stickiness
 
@@ -217,3 +201,13 @@ func TestGetDomainWithIncrementalrolloutOfGitLabSource(t *testing.T) {
 		})
 	}
 }
+
+func newTestDomains(t *testing.T, gitlabSource *MockSource) *Domains {
+	t.Helper()
+
+	return &Domains{
+		mu:     &sync.RWMutex{},
+		disk:   disk.New(),
+		gitlab: gitlabSource,
+	}
+}
diff --git a/internal/source/gitlab/api/client.go b/internal/source/gitlab/api/client.go
index 7206e25a..4a6c0893 100644
--- a/internal/source/gitlab/api/client.go
+++ b/internal/source/gitlab/api/client.go
@@ -2,10 +2,13 @@ package api
 
 import (
 	"context"
+	"time"
 )
 
 // Client represents an interface we use to retrieve information from GitLab
 type Client interface {
 	// GetLookup retrives an VirtualDomain from GitLab API and wraps it into Lookup
 	GetLookup(ctx context.Context, domain string) Lookup
+
+	Poll(retries int, interval time.Duration, errCh chan error)
 }
diff --git a/internal/source/gitlab/api/resolver.go b/internal/source/gitlab/api/resolver.go
index 061a1ddd..78cbc20b 100644
--- a/internal/source/gitlab/api/resolver.go
+++ b/internal/source/gitlab/api/resolver.go
@@ -2,11 +2,14 @@ package api
 
 import (
 	"context"
+	"time"
 )
 
 // Resolver represents an interface we use to retrieve information from GitLab
 // in a more generic way. It can be a concrete API client or cached client.
 type Resolver interface {
-	// Resolve retrives an VirtualDomain from GitLab API and wraps it into Lookup
+	// Resolve retrieves an VirtualDomain from GitLab API and wraps it into Lookup
 	Resolve(ctx context.Context, domain string) *Lookup
+	// Poll test
+	Poll(retries int, interval time.Duration, errCh chan error)
 }
diff --git a/internal/source/gitlab/cache/cache.go b/internal/source/gitlab/cache/cache.go
index c8d166b5..de91e8b6 100644
--- a/internal/source/gitlab/cache/cache.go
+++ b/internal/source/gitlab/cache/cache.go
@@ -109,3 +109,7 @@ func (c *Cache) Resolve(ctx context.Context, domain string) *api.Lookup {
 	metrics.DomainsSourceCacheMiss.Inc()
 	return entry.Retrieve(ctx, c.client)
 }
+
+func (c *Cache) Poll(retries int, interval time.Duration, errCh chan error) {
+	c.client.Poll(retries, interval, errCh)
+}
diff --git a/internal/source/gitlab/cache/cache_test.go b/internal/source/gitlab/cache/cache_test.go
index 7a12cd3b..27bd449a 100644
--- a/internal/source/gitlab/cache/cache_test.go
+++ b/internal/source/gitlab/cache/cache_test.go
@@ -69,6 +69,8 @@ func (c *client) GetLookup(ctx context.Context, _ string) api.Lookup {
 	return lookup
 }
 
+func (c *client) Poll(int, time.Duration, chan error) {}
+
 func withTestCache(config resolverConfig, cacheConfig *cacheConfig, block func(*Cache, *client)) {
 	var chanSize int
 
diff --git a/internal/source/gitlab/client/client_poll.go b/internal/source/gitlab/client/client_poll.go
index bdbfb345..cca7b731 100644
--- a/internal/source/gitlab/client/client_poll.go
+++ b/internal/source/gitlab/client/client_poll.go
@@ -3,16 +3,18 @@ package client
 import (
 	"fmt"
 	"time"
+
+	log "github.com/sirupsen/logrus"
 )
 
 const (
 	// DefaultPollingMaxRetries to be used by Poll
-	DefaultPollingMaxRetries = 30
+	DefaultPollingMaxRetries = 10
 	// DefaultPollingInterval to be used by Poll
 	DefaultPollingInterval = 10 * time.Second
 )
 
-// Poll tries to call the /internal/pages/status API endpoint for
+// Poll tries to call the /internal/pages/status API endpoint once plus
 // `retries` every `interval`.
 // TODO: should we consider using an exponential back-off approach?
 // https://pkg.go.dev/github.com/cenkalti/backoff/v4?tab=doc#pkg-examples
@@ -21,10 +23,14 @@ func (gc *Client) Poll(retries int, interval time.Duration, errCh chan error) {
 
 	var err error
 	for i := 0; i <= retries; i++ {
+		log.Info("polling GitLab internal pages status API")
 		err = gc.Status()
 		if err == nil {
+			log.Info("GitLab internal pages status API connected successfully")
+
 			// return as soon as we connect to the API
 			errCh <- nil
+			return
 		}
 
 		time.Sleep(interval)
diff --git a/internal/source/gitlab/client/client_poll_test.go b/internal/source/gitlab/client/client_poll_test.go
index 59da32c5..f92ac998 100644
--- a/internal/source/gitlab/client/client_poll_test.go
+++ b/internal/source/gitlab/client/client_poll_test.go
@@ -74,7 +74,6 @@ func TestClient_Poll(t *testing.T) {
 
 			go client.Poll(tt.retries, tt.interval, errCh)
 
-			// go func() {
 			select {
 			case err := <-errCh:
 				if tt.wantErr {
@@ -88,7 +87,6 @@ func TestClient_Poll(t *testing.T) {
 				t.Logf("%s timed out", tt.name)
 				t.FailNow()
 			}
-			// }()
 		})
 	}
 }
diff --git a/internal/source/gitlab/client/client_stub.go b/internal/source/gitlab/client/client_stub.go
index 604f127b..c3879a33 100644
--- a/internal/source/gitlab/client/client_stub.go
+++ b/internal/source/gitlab/client/client_stub.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/json"
 	"os"
+	"time"
 
 	"gitlab.com/gitlab-org/gitlab-pages/internal/source/gitlab/api"
 )
@@ -35,3 +36,7 @@ func (c StubClient) GetLookup(ctx context.Context, host string) api.Lookup {
 
 	return lookup
 }
+
+func (c StubClient) Poll(r int, i time.Duration, errCh chan error) {
+	errCh <- nil
+}
diff --git a/internal/source/gitlab/gitlab.go b/internal/source/gitlab/gitlab.go
index 12da9af1..c5095fb7 100644
--- a/internal/source/gitlab/gitlab.go
+++ b/internal/source/gitlab/gitlab.go
@@ -6,6 +6,7 @@ import (
 	"net/http"
 	"path"
 	"strings"
+	"time"
 
 	"gitlab.com/gitlab-org/gitlab-pages/internal/domain"
 	"gitlab.com/gitlab-org/gitlab-pages/internal/request"
@@ -94,3 +95,7 @@ func (g *Gitlab) Resolve(r *http.Request) (*serving.Request, error) {
 	return &serving.Request{Serving: defaultServing()},
 		errors.New("could not match lookup path")
 }
+
+func (g *Gitlab) Poll(retries int, interval time.Duration, errCh chan error) {
+	go g.client.Poll(retries, interval, errCh)
+}