Skip to content

Commit

Permalink
incusd/scriptlet: Add get_cluster_members
Browse files Browse the repository at this point in the history
Signed-off-by: Yueyuanmei Zhang <[email protected]>
  • Loading branch information
Yueyuanmei Zhang authored and stgraber committed Apr 30, 2024
1 parent 783b93e commit a1c34b1
Show file tree
Hide file tree
Showing 2 changed files with 90 additions and 0 deletions.
89 changes: 89 additions & 0 deletions internal/server/scriptlet/instance_placement.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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()
Expand Down
1 change: 1 addition & 0 deletions internal/server/scriptlet/load/load.go
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ func InstancePlacementCompile(src string) (*starlark.Program, error) {
"get_cluster_member_state",
"get_instance_resources",
"get_instances",
"get_cluster_members",
},
name)
}
Expand Down

0 comments on commit a1c34b1

Please sign in to comment.