Add NodesStatusInfo method to the Scylla client

There are places where we do not need State, but just Status (Up/Down).
This reduces calls to Scylla API in such cases from 5 to 2.

Author: czeslavo

pkg/probeserver/scylladbapistatus/prober.go

@@ -77,7 +77,7 @@ func (p *Prober) Readyz(w http.ResponseWriter, req *http.Request) {
 	defer scyllaClient.Close()
 
 	// Contact Scylla to learn about the status of the member
-	nodeStatuses, err := scyllaClient.Status(ctx, localhost)
+	nodeStatuses, err := scyllaClient.NodesStatusAndStateInfo(ctx, localhost)
 	if err != nil {
 		klog.ErrorS(err, "readyz probe: can't get scylla node status", "Service", p.serviceRef())
 		w.WriteHeader(http.StatusInternalServerError)

pkg/scyllaclient/client.go

@@ -85,30 +85,50 @@ func (c *Client) Close() {
 	}
 }
 
-func (c *Client) Status(ctx context.Context, host string) (NodeStatusInfoSlice, error) {
+func (c *Client) NodesStatusInfo(ctx context.Context, host string) (NodeStatusInfoSlice, error) {
 	if len(host) > 0 {
-		// Always query same host
 		ctx = forceHost(ctx, host)
 	}
 
-	// Get all hosts
-	resp, err := c.scyllaClient.Operations.StorageServiceHostIDGet(&scyllaoperations.StorageServiceHostIDGetParams{Context: ctx})
+	getHostIDResp, err := c.scyllaClient.Operations.StorageServiceHostIDGet(&scyllaoperations.StorageServiceHostIDGetParams{Context: ctx})
 	if err != nil {
 		return nil, err
 	}
 
-	all := make([]NodeStatusInfo, len(resp.Payload))
-	for i, p := range resp.Payload {
-		all[i].Addr = p.Key
-		all[i].HostID = p.Value
+	all := make([]NodeStatusInfo, len(getHostIDResp.Payload))
+	for i, n := range getHostIDResp.Payload {
+		all[i].Addr = n.Key
+		all[i].HostID = n.Value
+		all[i].Status = NodeStatusDown
 	}
 
-	// Get live nodes
-	live, err := c.scyllaClient.Operations.GossiperEndpointLiveGet(&scyllaoperations.GossiperEndpointLiveGetParams{Context: ctx})
+	getLiveNodesResp, err := c.scyllaClient.Operations.GossiperEndpointLiveGet(&scyllaoperations.GossiperEndpointLiveGetParams{Context: ctx})
 	if err != nil {
 		return nil, err
 	}
-	setNodeStatus(all, NodeStatusUp, live.Payload)
+	liveNodeAddrs := strset.New(getLiveNodesResp.Payload...)
+
+	for i, n := range all {
+		if liveNodeAddrs.Has(n.Addr) {
+			all[i].Status = NodeStatusUp
+		}
+	}
+
+	return all, nil
+}
+
+func (c *Client) NodesStatusAndStateInfo(ctx context.Context, host string) (NodeStatusAndStateInfoSlice, error) {
+	allNodesLivenessStatus, err := c.NodesStatusInfo(ctx, host)
+	if err != nil {
+		return nil, err
+	}
+
+	all := make([]NodeStatusAndStateInfo, len(allNodesLivenessStatus))
+	for i, n := range allNodesLivenessStatus {
+		all[i] = NodeStatusAndStateInfo{
+			NodeStatusInfo: n,
+		}
+	}
 
 	// Get joining nodes
 	joining, err := c.scyllaClient.Operations.StorageServiceNodesJoiningGet(&scyllaoperations.StorageServiceNodesJoiningGetParams{Context: ctx})
@@ -438,7 +458,7 @@ func DefaultTransport() *http.Transport {
 	}
 }
 
-func setNodeState(all []NodeStatusInfo, state NodeState, addrs []string) {
+func setNodeState(all []NodeStatusAndStateInfo, state NodeState, addrs []string) {
 	if len(addrs) == 0 {
 		return
 	}
@@ -451,19 +471,6 @@ func setNodeState(all []NodeStatusInfo, state NodeState, addrs []string) {
 	}
 }
 
-func setNodeStatus(all []NodeStatusInfo, status NodeStatus, addrs []string) {
-	if len(addrs) == 0 {
-		return
-	}
-	m := strset.New(addrs...)
-
-	for i := range all {
-		if m.Has(all[i].Addr) {
-			all[i].Status = status
-		}
-	}
-}
-
 // fixContentType adjusts Scylla REST API response so that it can be consumed
 // by Open API.
 func fixContentType(next http.RoundTripper) http.RoundTripper {

pkg/scyllaclient/model.go

@@ -127,29 +127,27 @@ const (
 	ReshapeCompactionType    CompactionType = "RESHAPE"
 )
 
-// NodeStatusInfo represents a nodetool status line.
-type NodeStatusInfo struct {
-	HostID string
-	Addr   string
-	Status NodeStatus
-	State  NodeState
+// NodeStatusAndStateInfo represents a node's status and state (like in nodetool status).
+type NodeStatusAndStateInfo struct {
+	NodeStatusInfo
+	State NodeState
 }
 
-func (s NodeStatusInfo) String() string {
+func (s NodeStatusAndStateInfo) String() string {
 	return fmt.Sprintf("host: %s, Status: %s%s", s.Addr, s.Status, s.State)
 }
 
 // IsUN returns true if host is Up and NORMAL meaning it's a fully functional
 // live node.
-func (s NodeStatusInfo) IsUN() bool {
+func (s NodeStatusAndStateInfo) IsUN() bool {
 	return s.Status == NodeStatusUp && s.State == NodeStateNormal
 }
 
-// NodeStatusInfoSlice adds functionality to Status response.
-type NodeStatusInfoSlice []NodeStatusInfo
+// NodeStatusAndStateInfoSlice adds functionality to Status response.
+type NodeStatusAndStateInfoSlice []NodeStatusAndStateInfo
 
 // Hosts returns slice of address of all nodes.
-func (s NodeStatusInfoSlice) Hosts() []string {
+func (s NodeStatusAndStateInfoSlice) Hosts() []string {
 	var hosts []string
 	for _, h := range s {
 		hosts = append(hosts, h.Addr)
@@ -158,7 +156,7 @@ func (s NodeStatusInfoSlice) Hosts() []string {
 }
 
 // HostIDs returns slice of HostID of all nodes.
-func (s NodeStatusInfoSlice) HostIDs() []string {
+func (s NodeStatusAndStateInfoSlice) HostIDs() []string {
 	var hostIDs []string
 	for _, h := range s {
 		hostIDs = append(hostIDs, h.HostID)
@@ -167,7 +165,7 @@ func (s NodeStatusInfoSlice) HostIDs() []string {
 }
 
 // LiveHosts returns slice of address of nodes in UN state.
-func (s NodeStatusInfoSlice) LiveHosts() []string {
+func (s NodeStatusAndStateInfoSlice) LiveHosts() []string {
 	var hosts []string
 	for _, h := range s {
 		if h.IsUN() {
@@ -178,7 +176,7 @@ func (s NodeStatusInfoSlice) LiveHosts() []string {
 }
 
 // DownHosts returns slice of address of nodes that are down.
-func (s NodeStatusInfoSlice) DownHosts() []string {
+func (s NodeStatusAndStateInfoSlice) DownHosts() []string {
 	var hosts []string
 	for _, h := range s {
 		if h.Status == NodeStatusDown {
@@ -189,7 +187,7 @@ func (s NodeStatusInfoSlice) DownHosts() []string {
 }
 
 // DownHostIDs returns slice of HostID of nodes that are down.
-func (s NodeStatusInfoSlice) DownHostIDs() []string {
+func (s NodeStatusAndStateInfoSlice) DownHostIDs() []string {
 	var hostIDs []string
 	for _, h := range s {
 		if h.Status == NodeStatusDown {
@@ -198,3 +196,43 @@ func (s NodeStatusInfoSlice) DownHostIDs() []string {
 	}
 	return hostIDs
 }
+
+// NodeStatusInfo represents the status (Up/Down) of a node.
+type NodeStatusInfo struct {
+	HostID string
+	Addr   string
+	Status NodeStatus
+}
+
+type NodeStatusInfoSlice []NodeStatusInfo
+
+// UpHostIDs returns slice of HostID of nodes that are up.
+func (s NodeStatusInfoSlice) UpHostIDs() []string {
+	var hosts []string
+	for _, h := range s {
+		if h.Status == NodeStatusUp {
+			hosts = append(hosts, h.Addr)
+		}
+	}
+	return hosts
+}
+
+// HostIDs returns slice of HostID of all nodes.
+func (s NodeStatusInfoSlice) HostIDs() []string {
+	var hostIDs []string
+	for _, h := range s {
+		hostIDs = append(hostIDs, h.HostID)
+	}
+	return hostIDs
+}
+
+// DownHostIDs returns slice of HostID of nodes that are down.
+func (s NodeStatusInfoSlice) DownHostIDs() []string {
+	var hosts []string
+	for _, h := range s {
+		if h.Status == NodeStatusDown {
+			hosts = append(hosts, h.HostID)
+		}
+	}
+	return hosts
+}

pkg/scyllaclient/status.go

@@ -3,29 +3,35 @@
 package scyllaclient
 
 import (
+	"errors"
+
 	"github.com/go-openapi/runtime"
-	"github.com/pkg/errors"
 	scyllav2models "github.com/scylladb/scylladb-swagger-go-client/scylladb/gen/v2/models"
 )
 
 // StatusCodeOf returns HTTP status code carried by the error or it's cause.
 // If not status can be found it returns 0.
 func StatusCodeOf(err error) int {
-	err = errors.Cause(err)
-	switch v := err.(type) {
-	case interface {
+	type coder interface {
 		Code() int
-	}:
-		return v.Code()
-	case *runtime.APIError:
-		return v.Code
-	case interface {
+	}
+	if v := new(coder); errors.As(err, v) {
+		return (*v).Code()
+	}
+
+	type payloader interface {
 		GetPayload() *scyllav2models.ErrorModel
-	}:
-		p := v.GetPayload()
+	}
+	if v := new(payloader); errors.As(err, v) {
+		p := (*v).GetPayload()
 		if p != nil {
 			return int(p.Code)
 		}
 	}
+
+	if v := new(runtime.APIError); errors.As(err, &v) {
+		return v.Code
+	}
+
 	return 0
 }

pkg/scyllaclient/status_test.go

@@ -0,0 +1,111 @@
+package scyllaclient_test
+
+import (
+	"errors"
+	"fmt"
+	"testing"
+
+	"github.com/go-openapi/runtime"
+	"github.com/scylladb/scylla-operator/pkg/scyllaclient"
+	scyllav2models "github.com/scylladb/scylladb-swagger-go-client/scylladb/gen/v2/models"
+)
+
+type mockCoderError struct {
+	code int
+}
+
+func (e *mockCoderError) Code() int {
+	return e.code
+}
+
+func (e *mockCoderError) Error() string {
+	return fmt.Sprintf("mock coder error with code %d", e.code)
+}
+
+type mockPayloaderError struct {
+	payload *scyllav2models.ErrorModel
+}
+
+func (e *mockPayloaderError) GetPayload() *scyllav2models.ErrorModel {
+	return e.payload
+}
+
+func (e *mockPayloaderError) Error() string {
+	if e.payload != nil {
+		return fmt.Sprintf("mock payloader error with code %d", e.payload.Code)
+	}
+	return "mock payloader error with nil payload"
+}
+
+func TestStatusCodeOf(t *testing.T) {
+	testCases := []struct {
+		name         string
+		err          error
+		expectedCode int
+	}{
+		{
+			name:         "coder interface - direct error",
+			err:          &mockCoderError{code: 404},
+			expectedCode: 404,
+		},
+		{
+			name:         "coder interface - wrapped error",
+			err:          fmt.Errorf("wrapped: %w", &mockCoderError{code: 500}),
+			expectedCode: 500,
+		},
+		{
+			name:         "payloader interface with valid payload - direct error",
+			err:          &mockPayloaderError{payload: &scyllav2models.ErrorModel{Code: 400}},
+			expectedCode: 400,
+		},
+		{
+			name:         "payloader interface with valid payload - wrapped error",
+			err:          fmt.Errorf("wrapped: %w", &mockPayloaderError{payload: &scyllav2models.ErrorModel{Code: 403}}),
+			expectedCode: 403,
+		},
+		{
+			name:         "payloader interface with nil payload - direct error",
+			err:          &mockPayloaderError{payload: nil},
+			expectedCode: 0,
+		},
+		{
+			name:         "payloader interface with nil payload - wrapped error",
+			err:          fmt.Errorf("wrapped: %w", &mockPayloaderError{payload: nil}),
+			expectedCode: 0,
+		},
+		{
+			name:         "runtime.APIError - direct error",
+			err:          &runtime.APIError{Code: 502},
+			expectedCode: 502,
+		},
+		{
+			name:         "runtime.APIError - wrapped error",
+			err:          fmt.Errorf("wrapped: %w", &runtime.APIError{Code: 503}),
+			expectedCode: 503,
+		},
+		{
+			name:         "unknown error type - direct error",
+			err:          errors.New("some unknown error"),
+			expectedCode: 0,
+		},
+		{
+			name:         "unknown error type - wrapped error",
+			err:          fmt.Errorf("wrapped: %w", errors.New("some unknown error")),
+			expectedCode: 0,
+		},
+		{
+			name:         "nil error",
+			err:          nil,
+			expectedCode: 0,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			result := scyllaclient.StatusCodeOf(tc.err)
+			if result != tc.expectedCode {
+				t.Errorf("StatusCodeOf(%v) = %d, expected %d", tc.err, result, tc.expectedCode)
+			}
+		})
+	}
+}

test/e2e/set/scyllacluster/scyllacluster_auth.go

@@ -182,12 +182,12 @@ var _ = g.Describe("ScyllaCluster authentication", func() {
 	})
 })
 
-func getScyllaClientStatus(ctx context.Context, hosts []string, authToken string) (scyllaclient.NodeStatusInfoSlice, error) {
+func getScyllaClientStatus(ctx context.Context, hosts []string, authToken string) (scyllaclient.NodeStatusAndStateInfoSlice, error) {
 	cfg := scyllaclient.DefaultConfig(authToken, hosts...)
 
 	client, err := scyllaclient.NewClient(cfg)
 	o.Expect(err).NotTo(o.HaveOccurred())
 	defer client.Close()
 
-	return client.Status(ctx, hosts[0])
+	return client.NodesStatusAndStateInfo(ctx, hosts[0])
 }