Skip to content

Commit 740b264

Browse files
mauriciocodermauriciocoder
authored
Merge branch 'main' into issues/1449
2 parents 1ef2a27 + 86edb93 commit 740b264

File tree

3 files changed

+122
-49
lines changed

3 files changed

+122
-49
lines changed

‎cmd/thv/app/stop.go

Lines changed: 120 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,15 @@
11
package app
22

33
import (
4+
"context"
45
"errors"
56
"fmt"
67

78
"github.com/spf13/cobra"
89
"golang.org/x/sync/errgroup"
910

1011
rt "github.com/stacklok/toolhive/pkg/container/runtime"
12+
"github.com/stacklok/toolhive/pkg/groups"
1113
"github.com/stacklok/toolhive/pkg/workloads"
1214
"github.com/stacklok/toolhive/pkg/workloads/types"
1315
)
@@ -24,25 +26,34 @@ var stopCmd = &cobra.Command{
2426
var (
2527
stopTimeout int
2628
stopAll bool
29+
stopGroup string
2730
)
2831

2932
func init() {
3033
stopCmd.Flags().IntVar(&stopTimeout, "timeout", 30, "Timeout in seconds before forcibly stopping the workload")
3134
stopCmd.Flags().BoolVar(&stopAll, "all", false, "Stop all running MCP servers")
35+
// TODO: Re-enable group flag when groups are implemented
36+
//stopCmd.Flags().StringVarP(&stopGroup, "group", "g", "", "Stop all MCP servers in a specific group")
37+
//
38+
//// Mark the flags as mutually exclusive
39+
//stopCmd.MarkFlagsMutuallyExclusive("all", "group")
40+
//
41+
//stopCmd.PreRunE = validateGroupFlag()
3242
}
3343

3444
// validateStopArgs validates the arguments for the stop command
3545
func validateStopArgs(cmd *cobra.Command, args []string) error {
36-
// Check if --all flag is set
46+
// Check if --all or --group flags are set
3747
all, _ := cmd.Flags().GetBool("all")
48+
group, _ := cmd.Flags().GetString("group")
3849

39-
if all {
40-
// If --all is set, no arguments should be provided
50+
if all || group != "" {
51+
// If --all or --group is set, no arguments should be provided
4152
if len(args) > 0 {
42-
return fmt.Errorf("no arguments should be provided when --all flag is set")
53+
return fmt.Errorf("no arguments should be provided when --all or --group flag is set")
4354
}
4455
} else {
45-
// If --all is not set, exactly one argument should be provided
56+
// If neither --all nor --group is set, exactly one argument should be provided
4657
if len(args) != 1 {
4758
return fmt.Errorf("exactly one workload name must be provided")
4859
}
@@ -61,59 +72,119 @@ func stopCmdFunc(cmd *cobra.Command, args []string) error {
6172

6273
var group *errgroup.Group
6374

64-
// Check if --all flag is set
6575
if stopAll {
66-
// Get list of all running workloads first
67-
workloadList, err := workloadManager.ListWorkloads(ctx, false) // false = only running workloads
68-
if err != nil {
69-
return fmt.Errorf("failed to list workloads: %v", err)
70-
}
76+
return stopAllWorkloads(ctx, workloadManager)
77+
}
7178

72-
// Extract workload names
73-
var workloadNames []string
74-
for _, workload := range workloadList {
75-
workloadNames = append(workloadNames, workload.Name)
76-
}
79+
if stopGroup != "" {
80+
return stopWorkloadsByGroup(ctx, workloadManager, stopGroup)
81+
}
7782

78-
if len(workloadNames) == 0 {
79-
fmt.Println("No running workloads to stop")
83+
// Stop single workload
84+
workloadName := args[0]
85+
group, err = workloadManager.StopWorkloads(ctx, []string{workloadName})
86+
if err != nil {
87+
// If the workload is not found or not running, treat as a non-fatal error.
88+
if errors.Is(err, rt.ErrWorkloadNotFound) ||
89+
errors.Is(err, workloads.ErrWorkloadNotRunning) ||
90+
errors.Is(err, types.ErrInvalidWorkloadName) {
91+
fmt.Printf("workload %s is not running\n", workloadName)
8092
return nil
8193
}
94+
return fmt.Errorf("unexpected error stopping workload: %v", err)
95+
}
8296

83-
// Stop all workloads using the bulk method
84-
group, err = workloadManager.StopWorkloads(ctx, workloadNames)
85-
if err != nil {
86-
return fmt.Errorf("failed to stop all workloads: %v", err)
87-
}
97+
// Since the stop operation is asynchronous, wait for the group to finish.
98+
if err := group.Wait(); err != nil {
99+
return fmt.Errorf("failed to stop workload %s: %v", workloadName, err)
100+
}
101+
fmt.Printf("workload %s stopped successfully\n", workloadName)
88102

89-
// Since the stop operation is asynchronous, wait for the group to finish.
90-
if err := group.Wait(); err != nil {
91-
return fmt.Errorf("failed to stop all workloads: %v", err)
92-
}
93-
fmt.Println("All workloads stopped successfully")
94-
} else {
95-
// Get workload name
96-
workloadName := args[0]
97-
98-
// Stop a single workload
99-
group, err = workloadManager.StopWorkloads(ctx, []string{workloadName})
100-
if err != nil {
101-
// If the workload is not found or not running, treat as a non-fatal error.
102-
if errors.Is(err, rt.ErrWorkloadNotFound) ||
103-
errors.Is(err, workloads.ErrWorkloadNotRunning) ||
104-
errors.Is(err, types.ErrInvalidWorkloadName) {
105-
fmt.Printf("workload %s is not running\n", workloadName)
106-
return nil
107-
}
108-
return fmt.Errorf("unexpected error stopping workload: %v", err)
109-
}
103+
return nil
104+
}
110105

111-
// Since the stop operation is asynchronous, wait for the group to finish.
112-
if err := group.Wait(); err != nil {
113-
return fmt.Errorf("failed to stop workload %s: %v", workloadName, err)
114-
}
115-
fmt.Printf("workload %s stopped successfully\n", workloadName)
106+
func stopAllWorkloads(ctx context.Context, workloadManager workloads.Manager) error {
107+
// Get list of all running workloads first
108+
workloadList, err := workloadManager.ListWorkloads(ctx, false) // false = only running workloads
109+
if err != nil {
110+
return fmt.Errorf("failed to list workloads: %v", err)
111+
}
112+
113+
// Extract workload names
114+
var workloadNames []string
115+
for _, workload := range workloadList {
116+
workloadNames = append(workloadNames, workload.Name)
117+
}
118+
119+
if len(workloadNames) == 0 {
120+
fmt.Println("No running workloads to stop")
121+
return nil
122+
}
123+
124+
// Stop all workloads using the bulk method
125+
group, err := workloadManager.StopWorkloads(ctx, workloadNames)
126+
if err != nil {
127+
return fmt.Errorf("failed to stop all workloads: %v", err)
128+
}
129+
130+
// Since the stop operation is asynchronous, wait for the group to finish.
131+
if err := group.Wait(); err != nil {
132+
return fmt.Errorf("failed to stop all workloads: %v", err)
133+
}
134+
fmt.Println("All workloads stopped successfully")
135+
return nil
136+
}
137+
138+
func stopWorkloadsByGroup(ctx context.Context, workloadManager workloads.Manager, groupName string) error {
139+
// Create a groups manager to list workloads in the group
140+
groupManager, err := groups.NewManager()
141+
if err != nil {
142+
return fmt.Errorf("failed to create group manager: %v", err)
143+
}
144+
145+
// Check if the group exists
146+
exists, err := groupManager.Exists(ctx, groupName)
147+
if err != nil {
148+
return fmt.Errorf("failed to check if group '%s' exists: %v", groupName, err)
149+
}
150+
if !exists {
151+
return fmt.Errorf("group '%s' does not exist", groupName)
152+
}
153+
154+
// Get list of running workloads and filter by group
155+
workloadList, err := workloadManager.ListWorkloads(ctx, false) // false = only running workloads
156+
if err != nil {
157+
return fmt.Errorf("failed to list running workloads: %v", err)
158+
}
159+
160+
// Filter workloads by group
161+
groupWorkloads, err := workloads.FilterByGroup(workloadList, groupName)
162+
if err != nil {
163+
return fmt.Errorf("failed to filter workloads by group: %v", err)
164+
}
165+
166+
if len(groupWorkloads) == 0 {
167+
fmt.Printf("No running MCP servers found in group '%s'\n", groupName)
168+
return nil
169+
}
170+
171+
// Extract workload names from the filtered list
172+
var workloadNames []string
173+
for _, workload := range groupWorkloads {
174+
workloadNames = append(workloadNames, workload.Name)
175+
}
176+
177+
// Stop workloads in the group
178+
subtasks, err := workloadManager.StopWorkloads(ctx, workloadNames)
179+
if err != nil {
180+
return fmt.Errorf("failed to stop workloads in group '%s': %v", groupName, err)
181+
}
182+
183+
// Wait for the stop operation to complete
184+
if err := subtasks.Wait(); err != nil {
185+
return fmt.Errorf("failed to stop workloads in group '%s': %v", groupName, err)
116186
}
117187

188+
fmt.Printf("Successfully stopped %d workload(s) in group '%s'\n", len(workloadNames), groupName)
118189
return nil
119190
}

‎go.mod

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ require (
2626
github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c
2727
github.com/prometheus/client_golang v1.23.0
2828
github.com/santhosh-tekuri/jsonschema/v5 v5.3.1
29+
github.com/santhosh-tekuri/jsonschema/v6 v6.0.2
2930
github.com/sigstore/protobuf-specs v0.5.0
3031
github.com/sigstore/sigstore-go v1.1.1
3132
github.com/spf13/viper v1.20.1

‎go.sum

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1463,6 +1463,7 @@ github.com/sagikazarmark/locafero v0.7.0 h1:5MqpDsTGNDhY8sGp0Aowyf0qKsPrhewaLSsF
14631463
github.com/sagikazarmark/locafero v0.7.0/go.mod h1:2za3Cg5rMaTMoG/2Ulr9AwtFaIppKXTRYnozin4aB5k=
14641464
github.com/santhosh-tekuri/jsonschema/v5 v5.3.1 h1:lZUw3E0/J3roVtGQ+SCrUrg3ON6NgVqpn3+iol9aGu4=
14651465
github.com/santhosh-tekuri/jsonschema/v5 v5.3.1/go.mod h1:uToXkOrWAZ6/Oc07xWQrPOhJotwFIyu2bBVN41fcDUY=
1466+
github.com/santhosh-tekuri/jsonschema/v6 v6.0.2/go.mod h1:JXeL+ps8p7/KNMjDQk3TCwPpBy0wYklyWTfbkIzdIFU=
14661467
github.com/sassoftware/relic v7.2.1+incompatible h1:Pwyh1F3I0r4clFJXkSI8bOyJINGqpgjJU3DYAZeI05A=
14671468
github.com/sassoftware/relic v7.2.1+incompatible/go.mod h1:CWfAxv73/iLZ17rbyhIEq3K9hs5w6FpNMdUT//qR+zk=
14681469
github.com/sassoftware/relic/v7 v7.6.2 h1:rS44Lbv9G9eXsukknS4mSjIAuuX+lMq/FnStgmZlUv4=

0 commit comments

Comments
 (0)