Expose GetServiceStatus function (#433)

* Expose GetServiceStatus function * service status rpc always return [serviceName].service
Snowflake-Labs · May 10, 2024 · c04c18f · c04c18f
1 parent 62be27c
commit c04c18f
Show file tree

Hide file tree

Showing 3 changed files with 47 additions and 34 deletions.
diff --git a/services/service/server/server_default.go b/services/service/server/server_default.go
@@ -20,9 +20,18 @@
 package server
 
 import (
+	"context"
+
 	pb "github.com/Snowflake-Labs/sansshell/services/service"
+	"github.com/coreos/go-systemd/v22/dbus"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 )
 
+func GetServiceStatus(ctx context.Context, conn *dbus.Conn, serviceName string, displayTimestamp bool) (*pb.StatusReply, error) {
+	return nil, status.Errorf(codes.Unimplemented, "check service status is not supported in OS other than linux")
+}
+
 func createServer() pb.ServiceServer {
 	return pb.UnimplementedServiceServer{}
 }
diff --git a/services/service/server/server_linux.go b/services/service/server/server_linux.go
@@ -227,37 +227,14 @@ func (s *server) List(ctx context.Context, req *pb.ListRequest) (*pb.ListReply,
 	return resp, nil
 }
 
-// See: pb.ServiceServer.Status
-func (s *server) Status(ctx context.Context, req *pb.StatusRequest) (*pb.StatusReply, error) {
+func GetServiceStatus(ctx context.Context, conn systemdConnection, serviceName string, displayTimestamp bool) (*pb.StatusReply, error) {
 	recorder := metrics.RecorderFromContextOrNoop(ctx)
 	logger := logr.FromContextOrDiscard(ctx)
-	if err := checkSupportedSystem(req.SystemType); err != nil {
-		return nil, err
-	}
-
-	unitName := req.GetServiceName()
-	if len(unitName) == 0 {
-		recorder.CounterOrLog(ctx, serviceStatusFailureCounter, 1, attribute.String("reason", "missing_name"))
-		return nil, status.Error(codes.InvalidArgument, "service name is required")
-	}
-
-	// Accept either 'foo' or 'foo.service'
-	if !strings.HasSuffix(unitName, unitSuffixService) {
-		unitName = unitName + unitSuffixService
-	}
-
-	conn, err := s.dialSystemd(ctx)
-	if err != nil {
-		recorder.CounterOrLog(ctx, serviceStatusFailureCounter, 1, attribute.String("reason", "dial_systemd_err"))
-		return nil, status.Errorf(codes.Internal, "error establishing systemd connection: %v", err)
-	}
-	defer conn.Close()
-
-	properties, err := conn.GetUnitPropertiesContext(ctx, unitName)
+	properties, err := conn.GetUnitPropertiesContext(ctx, serviceName)
 	if err != nil {
 		recorder.CounterOrLog(ctx, serviceStatusFailureCounter, 1, attribute.String("reason", "get_unit_properties_err"))
 		logger.V(3).Info("GetUnitPropertiesContext err: " + err.Error())
-		return nil, status.Errorf(codes.Internal, "failed to get unit properties of %s", req.GetServiceName())
+		return nil, status.Errorf(codes.Internal, "failed to get unit properties of %s", serviceName)
 	}
 	// cast map[string]interface{} to dbus.UnitStatus{} using json marshal + unmarshal
 	propertiesJson, err := json.Marshal(properties)
@@ -266,7 +243,6 @@ func (s *server) Status(ctx context.Context, req *pb.StatusRequest) (*pb.StatusR
 		logger.V(3).Info("failed to marshal properties: " + err.Error())
 		return nil, status.Errorf(codes.Internal, "failed to marshal unit properties to json")
 	}
-
 	unitState := dbus.UnitStatus{}
 	if errUnmarshal := json.Unmarshal(propertiesJson, &unitState); errUnmarshal != nil {
 		recorder.CounterOrLog(ctx, serviceStatusFailureCounter, 1, attribute.String("reason", "json_unmarshal_err"))
@@ -276,7 +252,7 @@ func (s *server) Status(ctx context.Context, req *pb.StatusRequest) (*pb.StatusR
 
 	unitStatus := unitStateToStatus(unitState)
 	var recentTimestampReachCurrentStatus *timestamppb.Timestamp
-	if req.DisplayTimestamp {
+	if displayTimestamp {
 		timestampReachCurrentStatus, err := getTimestampReachCurrentStatus(unitStatus, properties)
 		if err != nil {
 			recorder.CounterOrLog(ctx, serviceStatusFailureCounter, 1, attribute.String("reason", "get_timestamp_reach_current_status_err"))
@@ -288,13 +264,41 @@ func (s *server) Status(ctx context.Context, req *pb.StatusRequest) (*pb.StatusR
 	return &pb.StatusReply{
 		SystemType: pb.SystemType_SYSTEM_TYPE_SYSTEMD,
 		ServiceStatus: &pb.ServiceStatus{
-			ServiceName:                       req.GetServiceName(),
+			ServiceName:                       serviceName,
 			Status:                            unitStatus,
 			RecentTimestampReachCurrentStatus: recentTimestampReachCurrentStatus,
 		},
 	}, nil
 }
 
+// See: pb.ServiceServer.Status
+func (s *server) Status(ctx context.Context, req *pb.StatusRequest) (*pb.StatusReply, error) {
+	recorder := metrics.RecorderFromContextOrNoop(ctx)
+	if err := checkSupportedSystem(req.SystemType); err != nil {
+		return nil, err
+	}
+
+	unitName := req.GetServiceName()
+	if len(unitName) == 0 {
+		recorder.CounterOrLog(ctx, serviceStatusFailureCounter, 1, attribute.String("reason", "missing_name"))
+		return nil, status.Error(codes.InvalidArgument, "service name is required")
+	}
+
+	// Accept either 'foo' or 'foo.service'
+	if !strings.HasSuffix(unitName, unitSuffixService) {
+		unitName = unitName + unitSuffixService
+	}
+
+	conn, err := s.dialSystemd(ctx)
+	if err != nil {
+		recorder.CounterOrLog(ctx, serviceStatusFailureCounter, 1, attribute.String("reason", "dial_systemd_err"))
+		return nil, status.Errorf(codes.Internal, "error establishing systemd connection: %v", err)
+	}
+	defer conn.Close()
+
+	return GetServiceStatus(ctx, conn, unitName, req.DisplayTimestamp)
+}
+
 // See: pb.ServiceServer.Action
 func (s *server) Action(ctx context.Context, req *pb.ActionRequest) (*pb.ActionReply, error) {
 	recorder := metrics.RecorderFromContextOrNoop(ctx)

diff --git a/services/service/server/server_linux_test.go b/services/service/server/server_linux_test.go
@@ -299,7 +299,7 @@ func TestStatus(t *testing.T) {
 			want: &pb.StatusReply{
 				SystemType: pb.SystemType_SYSTEM_TYPE_SYSTEMD,
 				ServiceStatus: &pb.ServiceStatus{
-					ServiceName: "foo",
+					ServiceName: "foo.service",
 					Status:      pb.Status_STATUS_UNKNOWN,
 				},
 			},
@@ -365,7 +365,7 @@ func TestStatus(t *testing.T) {
 			want: &pb.StatusReply{
 				SystemType: pb.SystemType_SYSTEM_TYPE_SYSTEMD,
 				ServiceStatus: &pb.ServiceStatus{
-					ServiceName: "foo",
+					ServiceName: "foo.service",
 					Status:      pb.Status_STATUS_RUNNING,
 				},
 			},
@@ -393,7 +393,7 @@ func TestStatus(t *testing.T) {
 			want: &pb.StatusReply{
 				SystemType: pb.SystemType_SYSTEM_TYPE_SYSTEMD,
 				ServiceStatus: &pb.ServiceStatus{
-					ServiceName: "foo",
+					ServiceName: "foo.service",
 					Status:      pb.Status_STATUS_STOPPED,
 				},
 			},
@@ -417,7 +417,7 @@ func TestStatus(t *testing.T) {
 			want: &pb.StatusReply{
 				SystemType: pb.SystemType_SYSTEM_TYPE_SYSTEMD,
 				ServiceStatus: &pb.ServiceStatus{
-					ServiceName:                       "foo",
+					ServiceName:                       "foo.service",
 					Status:                            pb.Status_STATUS_RUNNING,
 					RecentTimestampReachCurrentStatus: timestamppb.New(time.UnixMicro(int64(activeEnterTimestamp))),
 				},
@@ -442,7 +442,7 @@ func TestStatus(t *testing.T) {
 			want: &pb.StatusReply{
 				SystemType: pb.SystemType_SYSTEM_TYPE_SYSTEMD,
 				ServiceStatus: &pb.ServiceStatus{
-					ServiceName:                       "foo",
+					ServiceName:                       "foo.service",
 					Status:                            pb.Status_STATUS_STOPPED,
 					RecentTimestampReachCurrentStatus: timestamppb.New(time.UnixMicro(int64(inactiveEnterTimestamp))),
 				},