diff --git a/mocks/tracker/announceclient/client.go b/mocks/tracker/announceclient/client.go index aa1fe9787..7c14f1947 100644 --- a/mocks/tracker/announceclient/client.go +++ b/mocks/tracker/announceclient/client.go @@ -1,5 +1,10 @@ // Code generated by MockGen. DO NOT EDIT. // Source: github.com/uber/kraken/tracker/announceclient (interfaces: Client) +// +// Generated by this command: +// +// mockgen -typed -package mockannounceclient . Client +// // Package mockannounceclient is a generated GoMock package. package mockannounceclient @@ -11,41 +16,103 @@ import ( time "time" ) -// MockClient is a mock of Client interface +// MockClient is a mock of Client interface. type MockClient struct { ctrl *gomock.Controller recorder *MockClientMockRecorder } -// MockClientMockRecorder is the mock recorder for MockClient +// MockClientMockRecorder is the mock recorder for MockClient. type MockClientMockRecorder struct { mock *MockClient } -// NewMockClient creates a new mock instance +// NewMockClient creates a new mock instance. func NewMockClient(ctrl *gomock.Controller) *MockClient { mock := &MockClient{ctrl: ctrl} mock.recorder = &MockClientMockRecorder{mock} return mock } -// EXPECT returns an object that allows the caller to indicate expected use +// EXPECT returns an object that allows the caller to indicate expected use. func (m *MockClient) EXPECT() *MockClientMockRecorder { return m.recorder } -// Announce mocks base method -func (m *MockClient) Announce(arg0 core.Digest, arg1 core.InfoHash, arg2 bool, arg3 int) ([]*core.PeerInfo, time.Duration, error) { +// Announce mocks base method. +func (m *MockClient) Announce(d core.Digest, h core.InfoHash, complete bool, version int) ([]*core.PeerInfo, time.Duration, error) { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "Announce", arg0, arg1, arg2, arg3) + ret := m.ctrl.Call(m, "Announce", d, h, complete, version) ret0, _ := ret[0].([]*core.PeerInfo) ret1, _ := ret[1].(time.Duration) ret2, _ := ret[2].(error) return ret0, ret1, ret2 } -// Announce indicates an expected call of Announce -func (mr *MockClientMockRecorder) Announce(arg0, arg1, arg2, arg3 interface{}) *gomock.Call { +// Announce indicates an expected call of Announce. +func (mr *MockClientMockRecorder) Announce(d, h, complete, version interface{}) *MockClientAnnounceCall { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Announce", reflect.TypeOf((*MockClient)(nil).Announce), arg0, arg1, arg2, arg3) + call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Announce", reflect.TypeOf((*MockClient)(nil).Announce), d, h, complete, version) + return &MockClientAnnounceCall{Call: call} +} + +// MockClientAnnounceCall wrap *gomock.Call +type MockClientAnnounceCall struct { + *gomock.Call +} + +// Return rewrite *gomock.Call.Return +func (c *MockClientAnnounceCall) Return(arg0 []*core.PeerInfo, arg1 time.Duration, arg2 error) *MockClientAnnounceCall { + c.Call = c.Call.Return(arg0, arg1, arg2) + return c +} + +// Do rewrite *gomock.Call.Do +func (c *MockClientAnnounceCall) Do(f func(core.Digest, core.InfoHash, bool, int) ([]*core.PeerInfo, time.Duration, error)) *MockClientAnnounceCall { + c.Call = c.Call.Do(f) + return c +} + +// DoAndReturn rewrite *gomock.Call.DoAndReturn +func (c *MockClientAnnounceCall) DoAndReturn(f func(core.Digest, core.InfoHash, bool, int) ([]*core.PeerInfo, time.Duration, error)) *MockClientAnnounceCall { + c.Call = c.Call.DoAndReturn(f) + return c +} + +// CheckReadiness mocks base method. +func (m *MockClient) CheckReadiness() error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "CheckReadiness") + ret0, _ := ret[0].(error) + return ret0 +} + +// CheckReadiness indicates an expected call of CheckReadiness. +func (mr *MockClientMockRecorder) CheckReadiness() *MockClientCheckReadinessCall { + mr.mock.ctrl.T.Helper() + call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CheckReadiness", reflect.TypeOf((*MockClient)(nil).CheckReadiness)) + return &MockClientCheckReadinessCall{Call: call} +} + +// MockClientCheckReadinessCall wrap *gomock.Call +type MockClientCheckReadinessCall struct { + *gomock.Call +} + +// Return rewrite *gomock.Call.Return +func (c *MockClientCheckReadinessCall) Return(arg0 error) *MockClientCheckReadinessCall { + c.Call = c.Call.Return(arg0) + return c +} + +// Do rewrite *gomock.Call.Do +func (c *MockClientCheckReadinessCall) Do(f func() error) *MockClientCheckReadinessCall { + c.Call = c.Call.Do(f) + return c +} + +// DoAndReturn rewrite *gomock.Call.DoAndReturn +func (c *MockClientCheckReadinessCall) DoAndReturn(f func() error) *MockClientCheckReadinessCall { + c.Call = c.Call.DoAndReturn(f) + return c } diff --git a/tracker/announceclient/client.go b/tracker/announceclient/client.go index 1085a29e8..94dfef126 100644 --- a/tracker/announceclient/client.go +++ b/tracker/announceclient/client.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, @@ -23,6 +23,7 @@ import ( "time" "github.com/uber/kraken/core" + "github.com/uber/kraken/lib/backend" "github.com/uber/kraken/lib/hashring" "github.com/uber/kraken/utils/httputil" ) @@ -58,6 +59,7 @@ type Response struct { // Client defines a client for announcing and getting peers. type Client interface { + CheckReadiness() error Announce( d core.Digest, h core.InfoHash, @@ -89,6 +91,18 @@ func getEndpoint(version int, addr string, h core.InfoHash) (method, url string) return "POST", fmt.Sprintf("http://%s/announce/%s", addr, h.String()) } +func (c *client) CheckReadiness() error { + addr := c.ring.Locations(backend.ReadinessCheckDigest)[0] + _, err := httputil.Get( + fmt.Sprintf("http://%s/readiness", addr), + httputil.SendTimeout(5*time.Second), + httputil.SendTLS(c.tls)) + if err != nil { + return fmt.Errorf("tracker not ready: %v", err) + } + return nil +} + // Announce announces the torrent identified by (d, h) with the number of // downloaded bytes. Returns a list of all other peers announcing for said torrent, // sorted by priority, and the interval for the next announce. @@ -142,6 +156,10 @@ func Disabled() Client { return DisabledClient{} } +func (c DisabledClient) CheckReadiness() error { + return nil +} + // Announce always returns error. func (c DisabledClient) Announce( d core.Digest, h core.InfoHash, complete bool, version int) ([]*core.PeerInfo, time.Duration, error) { diff --git a/tracker/trackerserver/announce_test.go b/tracker/trackerserver/announce_test.go index cc78f0cd3..7e28d05d8 100644 --- a/tracker/trackerserver/announce_test.go +++ b/tracker/trackerserver/announce_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, @@ -16,6 +16,7 @@ package trackerserver import ( "errors" "fmt" + "regexp" "testing" "time" @@ -33,6 +34,48 @@ func newAnnounceClient(pctx core.PeerContext, addr string) announceclient.Client return announceclient.New(pctx, hashring.NoopPassiveRing(hostlist.Fixture(addr)), nil) } +func TestCheckReadiness(t *testing.T) { + for _, tc := range []struct { + name string + mockOriginErr error + expectedErrMsgPattern string + }{ + { + name: "success", + mockOriginErr: nil, + expectedErrMsgPattern: "", + }, + { + name: "failure, 503 (origin fails)", + mockOriginErr: errors.New("origin error"), + expectedErrMsgPattern: fmt.Sprintf(`tracker not ready: GET http://127\.0\.0\.1:\d+/readiness 503: not ready to serve traffic: origin error`), + }, + } { + t.Run(tc.name, func(t *testing.T) { + require := require.New(t) + + mocks, cleanup := newServerMocks(t, Config{}) + defer cleanup() + + addr, stop := testutil.StartServer(mocks.handler()) + defer stop() + + mocks.originCluster.EXPECT().CheckReadiness().Return(tc.mockOriginErr) + + pctx := core.PeerContextFixture() + client := newAnnounceClient(pctx, addr) + + err := client.CheckReadiness() + if tc.expectedErrMsgPattern == "" { + require.Nil(err) + } else { + r, _ := regexp.Compile(tc.expectedErrMsgPattern) + require.True(r.MatchString(err.Error())) + } + }) + } +} + func TestAnnounceSinglePeerResponse(t *testing.T) { for _, version := range []int{announceclient.V1, announceclient.V2} { t.Run(fmt.Sprintf("V%d", version), func(t *testing.T) { diff --git a/tracker/trackerserver/server.go b/tracker/trackerserver/server.go index 7b98f4886..5637b4e28 100644 --- a/tracker/trackerserver/server.go +++ b/tracker/trackerserver/server.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, @@ -77,6 +77,8 @@ func (s *Server) Handler() http.Handler { r.Use(middleware.LatencyTimer(s.stats)) r.Get("/health", handler.Wrap(s.healthHandler)) + r.Get("/readiness", handler.Wrap(s.readinessCheckHandler)) + r.Get("/announce", handler.Wrap(s.announceHandlerV1)) r.Post("/announce/{infohash}", handler.Wrap(s.announceHandlerV2)) r.Get("/namespace/{namespace}/blobs/{digest}/metainfo", handler.Wrap(s.getMetaInfoHandler)) @@ -92,6 +94,15 @@ func (s *Server) ListenAndServe() error { return listener.Serve(s.config.Listener, s.Handler()) } +func (s *Server) readinessCheckHandler(w http.ResponseWriter, r *http.Request) error { + err := s.originCluster.CheckReadiness() + if err != nil { + return handler.Errorf("not ready to serve traffic: %s", err).Status(http.StatusServiceUnavailable) + } + fmt.Fprintln(w, "OK") + return nil +} + func (s *Server) healthHandler(w http.ResponseWriter, r *http.Request) error { fmt.Fprintln(w, "OK") return nil