Add ability to check backend readiness (#372)

As described in #371, all services should implement a "readiness" endpoint, which checks whether the service is ready to serve traffic, i.e. whether it can reach all of its dependencies. Thus, the origin and build index's readiness endpoints should check if they can reach remote backends. This diff adds a function to the backend manager that allows exactly that by iterating through the storage backends and running a Stat call on all of them.

I suggest it be configurable which backends need to be checked by the readiness endpoint (through the must_ready flag in the .yaml config), as some are more important for Kraken than others.

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/constants.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,6 +14,7 @@
 package backend
 
 import (
+	"github.com/uber/kraken/core"
 	"github.com/uber/kraken/utils/memsize"
 
 	"github.com/c2h5oh/datasize"
@@ -25,3 +26,8 @@ const (
 	DefaultConcurrency int               = 10
 	DefaultListMaxKeys int               = 250
 )
+
+var (
+	ReadinessCheckNamespace string = core.NamespaceFixture()
+	ReadinessCheckName      string = core.DigestFixture().Hex()
+)

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"
 
@@ -30,18 +31,20 @@ var (
 )
 
 type backend struct {
-	regexp *regexp.Regexp
-	client Client
+	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,
+		regexp:    re,
+		client:    c,
+		mustReady: mustReady,
 	}, nil
 }
 
@@ -80,7 +83,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 +115,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 +142,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) CheckReadiness() error {
+	for _, b := range m.backends {
+		if !b.mustReady {
+			continue
+		}
+		_, err := b.client.Stat(ReadinessCheckNamespace, ReadinessCheckName)
+		if err != nil && err != backenderrors.ErrBlobNotFound {
+			return fmt.Errorf("backend for namespace '%s' not ready: %s", b.regexp.String(), err)
+		}
+	}
+	return 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,14 @@
 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"
+	"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 +53,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 +69,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 +159,81 @@ func TestManagerBandwidth(t *testing.T) {
 
 	checkBandwidth(5, 25)
 }
+
+func TestManagerCheckReadiness(t *testing.T) {
+	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(backend.ReadinessCheckNamespace, backend.ReadinessCheckName).Return(mockStat1, tc.mockStat1Err).AnyTimes()
+			c2.EXPECT().Stat(backend.ReadinessCheckNamespace, backend.ReadinessCheckName).Return(mockStat2, tc.mockStat2Err).AnyTimes()
+
+			require.NoError(m.Register(n1, c1, tc.mustReady1))
+			require.NoError(m.Register(n2, c2, tc.mustReady2))
+
+			err := m.CheckReadiness()
+			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