More integration and testing of Store readiness

http/service.go

@@ -73,6 +73,9 @@ type Store interface {
 	// LeaderAddr returns the Raft address of the leader of the cluster.
 	LeaderAddr() (string, error)
 
+	// Ready returns whether the Store is ready to service requests.
+	Ready() bool
+
 	// Stats returns stats on the Store.
 	Stats() (map[string]interface{}, error)
 
@@ -1176,8 +1179,15 @@ func (s *Service) handleReadyz(w http.ResponseWriter, r *http.Request) {
 		w.Write([]byte(fmt.Sprintf("[+]node ok\n[+]leader not contactable: %s", err.Error())))
 		return
 	}
+
+	if !s.store.Ready() {
+		w.WriteHeader(http.StatusServiceUnavailable)
+		w.Write([]byte("[+]node ok\n[+]leader ok\n[+]store not ready"))
+		return
+	}
+
 	w.WriteHeader(http.StatusOK)
-	w.Write([]byte("[+]node ok\n[+]leader ok"))
+	w.Write([]byte("[+]node ok\n[+]leader ok\n[+]store ok"))
 }
 
 func (s *Service) handleExecute(w http.ResponseWriter, r *http.Request) {

http/service_test.go

@@ -1014,6 +1014,16 @@ func Test_Readyz(t *testing.T) {
 	if resp.StatusCode != http.StatusOK {
 		t.Fatalf("failed to get expected StatusOK for nodes, got %d", resp.StatusCode)
 	}
+
+	m.notReady = true
+	resp, err = client.Get(host + "/readyz")
+	if err != nil {
+		t.Fatalf("failed to make nodes request")
+	}
+	if resp.StatusCode != http.StatusServiceUnavailable {
+		t.Fatalf("failed to get expected StatusServiceUnavailable for nodes, got %d", resp.StatusCode)
+	}
+
 }
 
 func Test_ForwardingRedirectQuery(t *testing.T) {
@@ -1190,6 +1200,7 @@ type MockStore struct {
 	backupFn   func(br *command.BackupRequest, dst io.Writer) error
 	loadFn     func(lr *command.LoadRequest) error
 	leaderAddr string
+	notReady   bool // Default value is true, easier to test.
 }
 
 func (m *MockStore) Execute(er *command.ExecuteRequest) ([]*command.ExecuteResult, error) {
@@ -1222,6 +1233,10 @@ func (m *MockStore) LeaderAddr() (string, error) {
 	return m.leaderAddr, nil
 }
 
+func (m *MockStore) Ready() bool {
+	return !m.notReady
+}
+
 func (m *MockStore) Stats() (map[string]interface{}, error) {
 	return nil, nil
 }

store/store.go

@@ -456,14 +456,16 @@ func (s *Store) Ready() bool {
 		return false
 	}
 
-	s.readyChansMu.Lock()
-	defer s.readyChansMu.Unlock()
-	if s.numClosedReadyChannels != len(s.readyChans) {
-		return false
-	}
-	s.readyChans = nil
-	s.numClosedReadyChannels = 0
-	return true
+	return func() bool {
+		s.readyChansMu.Lock()
+		defer s.readyChansMu.Unlock()
+		if s.numClosedReadyChannels != len(s.readyChans) {
+			return false
+		}
+		s.readyChans = nil
+		s.numClosedReadyChannels = 0
+		return true
+	}()
 }
 
 // Close closes the store. If wait is true, waits for a graceful shutdown.

system_test/single_node_test.go

@@ -44,6 +44,19 @@ func Test_SingleNodeBasicEndpoint(t *testing.T) {
 	if !ready {
 		t.Fatalf("node is not ready")
 	}
+
+	// Test that registering a ready channel affects readiness as expected.
+	ch := make(chan struct{})
+	node.Store.RegisterReadyChannel(ch)
+	ready, err = node.Ready()
+	if err != nil {
+		t.Fatalf(`failed to retrieve readiness: %s`, err)
+	}
+	if ready {
+		t.Fatalf("node is ready when registered ready channel not yet closed")
+	}
+	close(ch)
+	testPoll(t, node.Ready, 100*time.Millisecond, 5*time.Second)
 }
 
 func Test_SingleNodeNotReadyLive(t *testing.T) {
@@ -1182,3 +1195,25 @@ func Test_SingleNodeNoopSnapLogsReopen(t *testing.T) {
 		}
 	}
 }
+
+func testPoll(t *testing.T, f func() (bool, error), p time.Duration, d time.Duration) {
+	tck := time.NewTicker(p)
+	defer tck.Stop()
+	tmr := time.NewTimer(d)
+	defer tmr.Stop()
+
+	for {
+		select {
+		case <-tck.C:
+			b, err := f()
+			if err != nil {
+				continue
+			}
+			if b {
+				return
+			}
+		case <-tmr.C:
+			t.Fatalf("timeout expired: %s", t.Name())
+		}
+	}
+}