kraken

1
// Copyright (c) 2016-2019 Uber Technologies, Inc.
2
//
3
// Licensed under the Apache License, Version 2.0 (the "License");
4
// you may not use this file except in compliance with the License.
5
// You may obtain a copy of the License at
6
//
7
//     http://www.apache.org/licenses/LICENSE-2.0
8
//
9
// Unless required by applicable law or agreed to in writing, software
10
// distributed under the License is distributed on an "AS IS" BASIS,
11
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
// See the License for the specific language governing permissions and
13
// limitations under the License.
14
package blobserver
15

16
import (
17
	"bytes"
18
	"testing"
19
	"time"
20

21
	"go.uber.org/zap"
22

23
	"github.com/andres-erbsen/clock"
24
	"github.com/golang/mock/gomock"
25
	"github.com/stretchr/testify/require"
26
	"github.com/uber-go/tally"
27

28
	"github.com/uber/kraken/core"
29
	"github.com/uber/kraken/lib/backend"
30
	"github.com/uber/kraken/lib/blobrefresh"
31
	"github.com/uber/kraken/lib/hashring"
32
	"github.com/uber/kraken/lib/healthcheck"
33
	"github.com/uber/kraken/lib/hostlist"
34
	"github.com/uber/kraken/lib/metainfogen"
35
	"github.com/uber/kraken/lib/store"
36
	"github.com/uber/kraken/mocks/lib/backend"
37
	"github.com/uber/kraken/mocks/lib/persistedretry"
38
	"github.com/uber/kraken/mocks/origin/blobclient"
39
	"github.com/uber/kraken/origin/blobclient"
40
	"github.com/uber/kraken/utils/log"
41
	"github.com/uber/kraken/utils/stringset"
42
	"github.com/uber/kraken/utils/testutil"
43
)
44

45
const (
46
	master1 = "dummy-origin-master01-zone2:80"
47
	master2 = "dummy-origin-master02-zone2:80"
48
	master3 = "dummy-origin-master03-zone2:80"
49
)
50

51
func init() {
52
	zapConfig := zap.NewProductionConfig()
53
	zapConfig.OutputPaths = []string{}
54
	log.ConfigureLogger(zapConfig)
55
}
56

57
func newHashRing(maxReplica int) hashring.Ring {
58
	return hashring.New(
59
		hashring.Config{MaxReplica: maxReplica},
60
		hostlist.Fixture(master1, master2, master3),
61
		healthcheck.IdentityFilter{})
62
}
63

64
func hashRingNoReplica() hashring.Ring   { return newHashRing(1) }
65
func hashRingSomeReplica() hashring.Ring { return newHashRing(2) }
66
func hashRingMaxReplica() hashring.Ring  { return newHashRing(3) }
67

68
// testClientProvider implements blobclient.ClientProvider. It maps origin hostnames to
69
// the local addresses they are running on, such that Provide("dummy-origin")
70
// can resolve a real address.
71
type testClientProvider struct {
72
	clients map[string]blobclient.Client
73
}
74

75
func newTestClientProvider() *testClientProvider {
76
	return &testClientProvider{make(map[string]blobclient.Client)}
77
}
78

79
func (p *testClientProvider) register(host string, client blobclient.Client) {
80
	p.clients[host] = client
81
}
82

83
func (p *testClientProvider) Provide(host string) blobclient.Client {
84
	c, ok := p.clients[host]
85
	if !ok {
86
		log.Panicf("host %q not found", host)
87
	}
88
	return c
89
}
90

91
// testServer is a convenience wrapper around the underlying components of a
92
// Server and faciliates restarting Servers with new configuration.
93
type testServer struct {
94
	ctrl             *gomock.Controller
95
	host             string
96
	addr             string
97
	cas              *store.CAStore
98
	cp               *testClientProvider
99
	clusterProvider  *mockblobclient.MockClusterProvider
100
	pctx             core.PeerContext
101
	backendManager   *backend.Manager
102
	writeBackManager *mockpersistedretry.MockManager
103
	clk              *clock.Mock
104
	cleanup          func()
105
}
106

107
func newTestServer(
108
	t *testing.T, host string, ring hashring.Ring, cp *testClientProvider) *testServer {
109

110
	var cleanup testutil.Cleanup
111
	defer cleanup.Recover()
112

113
	ctrl := gomock.NewController(t)
114
	cleanup.Add(ctrl.Finish)
115

116
	clusterProvider := mockblobclient.NewMockClusterProvider(ctrl)
117

118
	pctx := core.PeerContextFixture()
119

120
	cas, c := store.CAStoreFixture()
121
	cleanup.Add(c)
122

123
	bm := backend.ManagerFixture()
124

125
	writeBackManager := mockpersistedretry.NewMockManager(ctrl)
126

127
	mg := metainfogen.Fixture(cas, 4)
128

129
	br := blobrefresh.New(blobrefresh.Config{}, tally.NoopScope, cas, bm, mg)
130

131
	clk := clock.NewMock()
132
	clk.Set(time.Now())
133

134
	s, err := New(
135
		Config{}, tally.NoopScope, clk, host, ring, cas, cp, clusterProvider, pctx,
136
		bm, br, mg, writeBackManager)
137
	if err != nil {
138
		panic(err)
139
	}
140

141
	addr, stop := testutil.StartServer(s.Handler())
142
	cleanup.Add(stop)
143

144
	cp.register(host, blobclient.New(addr, blobclient.WithChunkSize(16)))
145

146
	return &testServer{
147
		ctrl:             ctrl,
148
		host:             host,
149
		addr:             addr,
150
		cas:              cas,
151
		cp:               cp,
152
		clusterProvider:  clusterProvider,
153
		pctx:             pctx,
154
		backendManager:   bm,
155
		writeBackManager: writeBackManager,
156
		clk:              clk,
157
		cleanup:          cleanup.Run,
158
	}
159
}
160

161
func (s *testServer) backendClient(namespace string) *mockbackend.MockClient {
162
	client := mockbackend.NewMockClient(s.ctrl)
163
	if err := s.backendManager.Register(namespace, client); err != nil {
164
		panic(err)
165
	}
166
	return client
167
}
168

169
func (s *testServer) expectRemoteCluster(dns string) *mockblobclient.MockClusterClient {
170
	cc := mockblobclient.NewMockClusterClient(s.ctrl)
171
	s.clusterProvider.EXPECT().Provide(dns).Return(cc, nil).MinTimes(1)
172
	return cc
173
}
174

175
// computeBlobForHosts generates a random digest / content which shards to hosts.
176
func computeBlobForHosts(ring hashring.Ring, hosts ...string) *core.BlobFixture {
177
	want := stringset.New(hosts...)
178
	for {
179
		blob := core.SizedBlobFixture(32, 4)
180
		got := stringset.New(ring.Locations(blob.Digest)...)
181
		if stringset.Equal(want, got) {
182
			return blob
183
		}
184
	}
185
}
186

187
func ensureHasBlob(t *testing.T, c blobclient.Client, namespace string, blob *core.BlobFixture) {
188
	var buf bytes.Buffer
189
	require.NoError(t, c.DownloadBlob(namespace, blob.Digest, &buf))
190
	require.Equal(t, string(blob.Content), buf.String())
191
}
192