From dbfa267fd54f11bf613a4fc2d973857e81cfd64f Mon Sep 17 00:00:00 2001
From: Yueyuanmei Zhang <yymzhang@cs.utexas.edu>
Date: Mon, 22 Apr 2024 15:06:08 -0500
Subject: [PATCH] incusd/scriptlet: Add get_cluster_members

Signed-off-by: Yueyuanmei Zhang <yymzhang@cs.utexas.edu>
---
 .../server/scriptlet/instance_placement.go    | 89 +++++++++++++++++++
 internal/server/scriptlet/load/load.go        |  1 +
 2 files changed, 90 insertions(+)

diff --git a/internal/server/scriptlet/instance_placement.go b/internal/server/scriptlet/instance_placement.go
index 2c29394288f..beaaefdaa17 100644
--- a/internal/server/scriptlet/instance_placement.go
+++ b/internal/server/scriptlet/instance_placement.go
@@ -294,6 +294,94 @@ func InstancePlacementRun(ctx context.Context, l logger.Logger, s *state.State,
 		return rv, nil
 	}
 
+	getClusterMembersFunc := func(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
+		var group string
+		var allMembers []db.NodeInfo
+
+		err := starlark.UnpackArgs(b.Name(), args, kwargs, "group", &group)
+		if err != nil {
+			return nil, err
+		}
+
+		err = s.DB.Cluster.Transaction(ctx, func(ctx context.Context, tx *db.ClusterTx) error {
+			allMembers, err = tx.GetNodes(ctx)
+			if err != nil {
+				return err
+			}
+
+			allMembers, err = tx.GetCandidateMembers(ctx, allMembers, nil, group, nil, s.GlobalConfig.OfflineThreshold())
+			if err != nil {
+				return err
+			}
+
+			return nil
+		})
+		if err != nil {
+			return nil, err
+		}
+
+		var raftNodes []db.RaftNode
+		err = s.DB.Node.Transaction(ctx, func(ctx context.Context, tx *db.NodeTx) error {
+			raftNodes, err = tx.GetRaftNodes(ctx)
+			if err != nil {
+				return fmt.Errorf("Failed loading RAFT nodes: %w", err)
+			}
+
+			return nil
+		})
+		if err != nil {
+			return nil, err
+		}
+
+		allMembersInfo := make([]*api.ClusterMember, 0, len(allMembers))
+		err = s.DB.Cluster.Transaction(ctx, func(ctx context.Context, tx *db.ClusterTx) error {
+			failureDomains, err := tx.GetFailureDomainsNames(ctx)
+			if err != nil {
+				return fmt.Errorf("Failed loading failure domains names: %w", err)
+			}
+
+			memberFailureDomains, err := tx.GetNodesFailureDomains(ctx)
+			if err != nil {
+				return fmt.Errorf("Failed loading member failure domains: %w", err)
+			}
+
+			maxVersion, err := tx.GetNodeMaxVersion(ctx)
+			if err != nil {
+				return fmt.Errorf("Failed getting max member version: %w", err)
+			}
+
+			args := db.NodeInfoArgs{
+				LeaderAddress:        leaderAddress,
+				FailureDomains:       failureDomains,
+				MemberFailureDomains: memberFailureDomains,
+				OfflineThreshold:     s.GlobalConfig.OfflineThreshold(),
+				MaxMemberVersion:     maxVersion,
+				RaftNodes:            raftNodes,
+			}
+
+			for i := range allMembers {
+				candidateMemberInfo, err := allMembers[i].ToAPI(ctx, tx, args)
+				if err != nil {
+					return err
+				}
+
+				allMembersInfo = append(allMembersInfo, candidateMemberInfo)
+			}
+
+			return nil
+		})
+		if err != nil {
+			return nil, err
+		}
+
+		rv, err := StarlarkMarshal(allMembersInfo)
+		if err != nil {
+			return nil, fmt.Errorf("Marshalling instance resources failed: %w", err)
+		}
+
+		return rv, nil
+	}
+
 	var err error
 	var raftNodes []db.RaftNode
 	err = s.DB.Node.Transaction(ctx, func(ctx context.Context, tx *db.NodeTx) error {
@@ -360,6 +448,7 @@ func InstancePlacementRun(ctx context.Context, l logger.Logger, s *state.State,
 		"get_cluster_member_state":     starlark.NewBuiltin("get_cluster_member_state", getClusterMemberStateFunc),
 		"get_instance_resources":       starlark.NewBuiltin("get_instance_resources", getInstanceResourcesFunc),
 		"get_instances":                starlark.NewBuiltin("get_instances", getInstancesFunc),
+		"get_cluster_members":          starlark.NewBuiltin("get_cluster_members", getClusterMembersFunc),
 	}
 
 	prog, thread, err := scriptletLoad.InstancePlacementProgram()
diff --git a/internal/server/scriptlet/load/load.go b/internal/server/scriptlet/load/load.go
index 2556230196f..894c16830ad 100644
--- a/internal/server/scriptlet/load/load.go
+++ b/internal/server/scriptlet/load/load.go
@@ -23,6 +23,7 @@ func InstancePlacementCompile(src string) (*starlark.Program, error) {
 			"get_cluster_member_state",
 			"get_instance_resources",
 			"get_instances",
+			"get_cluster_members",
 		},
 			name)
 	}