Add ability to check backend readiness

build-index/tagserver/server_test.go

@@ -76,7 +76,7 @@ func newServerMocks(t *testing.T) (*serverMocks, func()) {
 
 	backends := backend.ManagerFixture()
 	backendClient := mockbackend.NewMockClient(ctrl)
-	require.NoError(t, backends.Register(_testNamespace, backendClient))
+	require.NoError(t, backends.Register(_testNamespace, backendClient, false))
 
 	remotes, err := tagreplication.RemotesConfig{
 		_testRemote: []string{_testNamespace},

build-index/tagstore/store_test.go

@@ -57,7 +57,7 @@ func newStoreMocks(t *testing.T) (*storeMocks, func()) {
 
 	backends := backend.ManagerFixture()
 	backendClient := mockbackend.NewMockClient(ctrl)
-	require.NoError(t, backends.Register(_testNamespace, backendClient))
+	require.NoError(t, backends.Register(_testNamespace, backendClient, false))
 
 	writeBackManager := mockpersistedretry.NewMockManager(ctrl)
 

lib/backend/config.go

@@ -4,7 +4,7 @@
 // you may not use this file except in compliance with the License.
 // You may obtain a copy of the License at
 //
-//     http://www.apache.org/licenses/LICENSE-2.0
+//	http://www.apache.org/licenses/LICENSE-2.0
 //
 // Unless required by applicable law or agreed to in writing, software
 // distributed under the License is distributed on an "AS IS" BASIS,
@@ -26,12 +26,14 @@ type Config struct {
 
 	// If enabled, throttles upload / download bandwidth.
 	Bandwidth bandwidth.Config `yaml:"bandwidth"`
+	// Whether the service readiness endpoint will check the backend's readiness.
+	MustReady bool             `yaml:"must_ready"`
 }
 
 func (c Config) applyDefaults() Config {
 	for k := range c.Backend {
 		// TODO: don't hard code backend client name
-		if k == "s3" || k == "gcs"{
+		if k == "s3" || k == "gcs" {
 			if c.Bandwidth.IngressBitsPerSec == 0 {
 				c.Bandwidth.IngressBitsPerSec = 10 * 8 * memsize.Gbit
 			}

lib/backend/manager.go

@@ -18,6 +18,7 @@ import (
 	"fmt"
 	"regexp"
 
+	"github.com/uber/kraken/lib/backend/backenderrors"
 	"github.com/uber/kraken/utils/bandwidth"
 	"github.com/uber/kraken/utils/log"
 
@@ -28,20 +29,24 @@ import (
 var (
 	ErrNamespaceNotFound = errors.New("no matches for namespace")
 )
+const isReadyNamespace = "isReadyNamespace"
+const isReadyName = "38a03d499119bc417b8a6a016f2cb4540b9f9cc0c13e4da42a73867120d3e908"
 
 type backend struct {
 	regexp *regexp.Regexp
 	client Client
+	mustReady bool
 }
 
-func newBackend(namespace string, c Client) (*backend, error) {
+func newBackend(namespace string, c Client, mustReady bool) (*backend, error) {
 	re, err := regexp.Compile(namespace)
 	if err != nil {
 		return nil, fmt.Errorf("regexp: %s", err)
 	}
 	return &backend{
 		regexp: re,
 		client: c,
+		mustReady: mustReady,
 	}, nil
 }
 
@@ -80,7 +85,7 @@ func NewManager(configs []Config, auth AuthConfig, stats tally.Scope) (*Manager,
 			}
 			c = throttle(c, l)
 		}
-		b, err := newBackend(config.Namespace, c)
+		b, err := newBackend(config.Namespace, c, config.MustReady)
 		if err != nil {
 			return nil, fmt.Errorf("new backend for namespace %s: %s", config.Namespace, err)
 		}
@@ -112,13 +117,13 @@ func (m *Manager) AdjustBandwidth(denominator int) error {
 // Register dynamically registers a namespace with a provided client. Register
 // should be primarily used for testing purposes -- normally, namespaces should
 // be statically configured and provided upon construction of the Manager.
-func (m *Manager) Register(namespace string, c Client) error {
+func (m *Manager) Register(namespace string, c Client, mustReady bool) error {
 	for _, b := range m.backends {
 		if b.regexp.String() == namespace {
 			return fmt.Errorf("namespace %s already exists", namespace)
 		}
 	}
-	b, err := newBackend(namespace, c)
+	b, err := newBackend(namespace, c, mustReady)
 	if err != nil {
 		return fmt.Errorf("new backend: %s", err)
 	}
@@ -139,3 +144,18 @@ func (m *Manager) GetClient(namespace string) (Client, error) {
 	}
 	return nil, ErrNamespaceNotFound
 }
+
+// IsReady returns whether the backends are ready (reachable).
+// A backend must be explicitly configured as required for readiness to be checked.
+func (m *Manager) IsReady() (bool, error) {
+	for _, b := range m.backends {
+		if !b.mustReady {
+			continue
+		}
+		_, err := b.client.Stat(isReadyNamespace, isReadyName)
+		if err != nil && err != backenderrors.ErrBlobNotFound {
+			return false, fmt.Errorf("backend for namespace %s not ready: %s", b.regexp.String(), err)
+		}
+	}
+	return true, nil
+}

lib/backend/manager_test.go

@@ -4,7 +4,7 @@
 // you may not use this file except in compliance with the License.
 // You may obtain a copy of the License at
 //
-//     http://www.apache.org/licenses/LICENSE-2.0
+//	http://www.apache.org/licenses/LICENSE-2.0
 //
 // Unless required by applicable law or agreed to in writing, software
 // distributed under the License is distributed on an "AS IS" BASIS,
@@ -14,10 +14,13 @@
 package backend_test
 
 import (
+	"errors"
 	"testing"
 
 	"github.com/uber-go/tally"
+	"github.com/uber/kraken/core"
 	. "github.com/uber/kraken/lib/backend"
+	"github.com/uber/kraken/lib/backend/backenderrors"
 	"github.com/uber/kraken/lib/backend/namepath"
 	"github.com/uber/kraken/lib/backend/testfs"
 	"github.com/uber/kraken/mocks/lib/backend"
@@ -49,8 +52,8 @@ func TestManagerNamespaceMatching(t *testing.T) {
 
 			m := ManagerFixture()
 
-			require.NoError(m.Register("static", c1))
-			require.NoError(m.Register("namespace-foo/.*", c2))
+			require.NoError(m.Register("static", c1, false))
+			require.NoError(m.Register("namespace-foo/.*", c2, false))
 
 			result, err := m.GetClient(test.namespace)
 			require.NoError(err)
@@ -65,8 +68,8 @@ func TestManagerErrDuplicateNamespace(t *testing.T) {
 
 	c := &NoopClient{}
 	m := ManagerFixture()
-	require.NoError(m.Register("static", c))
-	require.Error(m.Register("static", c))
+	require.NoError(m.Register("static", c, false))
+	require.Error(m.Register("static", c, false))
 }
 
 func TestManagerErrNamespaceNotFound(t *testing.T) {
@@ -155,3 +158,85 @@ func TestManagerBandwidth(t *testing.T) {
 
 	checkBandwidth(5, 25)
 }
+
+func TestManagerIsReady(t *testing.T) {
+	const isReadyNamespace = "isReadyNamespace"
+	const isReadyName = "38a03d499119bc417b8a6a016f2cb4540b9f9cc0c13e4da42a73867120d3e908"
+
+	n1 := "foo/*"
+	n2 := "bar/*"
+
+	tests := []struct {
+		name         string
+		mustReady1   bool
+		mustReady2   bool
+		mockStat1Err error
+		mockStat2Err error
+		expectedRes  bool
+		expectedErr  error
+	}{
+		{
+			name:         "both required, both pass (one with nil, one with NotFound)",
+			mustReady1:   true,
+			mustReady2:   true,
+			mockStat1Err: nil,
+			mockStat2Err: backenderrors.ErrBlobNotFound,
+			expectedRes:  true,
+			expectedErr:  nil,
+		},
+		{
+			name:         "both required, only second fails",
+			mustReady1:   true,
+			mustReady2:   true,
+			mockStat1Err: nil,
+			mockStat2Err: errors.New("network error"),
+			expectedRes:  false,
+			expectedErr:  errors.New("backend for namespace bar/* not ready: network error"),
+		},
+		{
+			name:         "second required, only first fails",
+			mustReady1:   false,
+			mustReady2:   true,
+			mockStat1Err: errors.New("network error"),
+			mockStat2Err: backenderrors.ErrBlobNotFound,
+			expectedRes:  true,
+			expectedErr:  nil,
+		},
+	}
+
+	for _, tc := range tests {
+		ctrl := gomock.NewController(t)
+		defer ctrl.Finish()
+		c1 := mockbackend.NewMockClient(ctrl)
+		c2 := mockbackend.NewMockClient(ctrl)
+
+		t.Run(tc.name, func(t *testing.T) {
+			require := require.New(t)
+
+			m := ManagerFixture()
+
+			mockStat1 := &core.BlobInfo{}
+			mockStat2 := &core.BlobInfo{}
+			if tc.mockStat1Err != nil {
+				mockStat1 = nil
+			}
+			if tc.mockStat2Err != nil {
+				mockStat2 = nil
+			}
+
+			c1.EXPECT().Stat(isReadyNamespace, isReadyName).Return(mockStat1, tc.mockStat1Err).AnyTimes()
+			c2.EXPECT().Stat(isReadyNamespace, isReadyName).Return(mockStat2, tc.mockStat2Err).AnyTimes()
+
+			require.NoError(m.Register(n1, c1, tc.mustReady1))
+			require.NoError(m.Register(n2, c2, tc.mustReady2))
+
+			res, err := m.IsReady()
+			require.Equal(tc.expectedRes, res)
+			if tc.expectedErr == nil {
+				require.NoError(err)
+			} else {
+				require.Equal(tc.expectedErr, err)
+			}
+		})
+	}
+}

lib/blobrefresh/refresher_test.go

@@ -63,7 +63,7 @@ func (m *refresherMocks) new() *Refresher {
 
 func (m *refresherMocks) newClient(namespace string) *mockbackend.MockClient {
 	client := mockbackend.NewMockClient(m.ctrl)
-	m.backends.Register(namespace, client)
+	m.backends.Register(namespace, client, false)
 	return client
 }
 

lib/persistedretry/writeback/executor_test.go

@@ -61,7 +61,7 @@ func (m *executorMocks) new() *Executor {
 
 func (m *executorMocks) client(namespace string) *mockbackend.MockClient {
 	client := mockbackend.NewMockClient(m.ctrl)
-	if err := m.backends.Register(namespace, client); err != nil {
+	if err := m.backends.Register(namespace, client, false); err != nil {
 		panic(err)
 	}
 	return client

lib/torrent/storage/originstorage/torrent_archive_test.go

@@ -52,7 +52,7 @@ func newArchiveMocks(t *testing.T, namespace string) (*archiveMocks, func()) {
 
 	backendClient := mockbackend.NewMockClient(ctrl)
 	backends := backend.ManagerFixture()
-	backends.Register(namespace, backendClient)
+	backends.Register(namespace, backendClient, false)
 
 	blobRefresher := blobrefresh.New(
 		blobrefresh.Config{}, tally.NoopScope, cas, backends, metainfogen.Fixture(cas, pieceLength))

origin/blobserver/testutils_test.go

@@ -160,7 +160,7 @@ func newTestServer(
 
 func (s *testServer) backendClient(namespace string) *mockbackend.MockClient {
 	client := mockbackend.NewMockClient(s.ctrl)
-	if err := s.backendManager.Register(namespace, client); err != nil {
+	if err := s.backendManager.Register(namespace, client, false); err != nil {
 		panic(err)
 	}
 	return client