Skip to content

Commit e7c2281

Browse files
kylecarbskylecarbs
committed
feat: Add "state" command to pull and push workspace state
It's possible for a workspace to become in an invalid state. This is something we'll detect for jobs, and allow monitoring of. These commands will allow admins to manually reconcile state.
1 parent 43c6bff commit e7c2281

File tree

8 files changed

+370
-1
lines changed

8 files changed

+370
-1
lines changed

‎cli/root.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,7 @@ func Root() *cobra.Command {
8080
server(),
8181
show(),
8282
start(),
83+
state(),
8384
stop(),
8485
ssh(),
8586
templates(),

‎cli/state.go

Lines changed: 111 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,111 @@
1+
package cli
2+
3+
import (
4+
"io"
5+
"os"
6+
7+
"github.com/spf13/cobra"
8+
9+
"github.com/coder/coder/codersdk"
10+
)
11+
12+
func state() *cobra.Command {
13+
cmd := &cobra.Command{
14+
Use: "state",
15+
}
16+
cmd.AddCommand(statePull(), statePush())
17+
return cmd
18+
}
19+
20+
func statePull() *cobra.Command {
21+
var buildName string
22+
cmd := &cobra.Command{
23+
Use: "pull <workspace> [file]",
24+
RunE: func(cmd *cobra.Command, args []string) error {
25+
client, err := createClient(cmd)
26+
if err != nil {
27+
return err
28+
}
29+
organization, err := currentOrganization(cmd, client)
30+
if err != nil {
31+
return err
32+
}
33+
34+
workspace, err := client.WorkspaceByOwnerAndName(cmd.Context(), organization.ID, codersdk.Me, args[0])
35+
if err != nil {
36+
return err
37+
}
38+
var build codersdk.WorkspaceBuild
39+
if buildName == "latest" {
40+
build = workspace.LatestBuild
41+
} else {
42+
build, err = client.WorkspaceBuildByName(cmd.Context(), workspace.ID, buildName)
43+
if err != nil {
44+
return err
45+
}
46+
}
47+
48+
state, err := client.WorkspaceBuildState(cmd.Context(), build.ID)
49+
if err != nil {
50+
return err
51+
}
52+
53+
if len(args) < 2 {
54+
cmd.Println(string(state))
55+
return nil
56+
}
57+
58+
return os.WriteFile(args[1], state, 0600)
59+
},
60+
}
61+
cmd.Flags().StringVarP(&buildName, "build", "b", "latest", "Specify a workspace build to target by name.")
62+
return cmd
63+
}
64+
65+
func statePush() *cobra.Command {
66+
var buildName string
67+
cmd := &cobra.Command{
68+
Use: "push <workspace> <file>",
69+
Args: cobra.ExactArgs(2),
70+
RunE: func(cmd *cobra.Command, args []string) error {
71+
client, err := createClient(cmd)
72+
if err != nil {
73+
return err
74+
}
75+
organization, err := currentOrganization(cmd, client)
76+
if err != nil {
77+
return err
78+
}
79+
80+
workspace, err := client.WorkspaceByOwnerAndName(cmd.Context(), organization.ID, codersdk.Me, args[0])
81+
if err != nil {
82+
return err
83+
}
84+
var build codersdk.WorkspaceBuild
85+
if buildName == "latest" {
86+
build = workspace.LatestBuild
87+
} else {
88+
build, err = client.WorkspaceBuildByName(cmd.Context(), workspace.ID, buildName)
89+
if err != nil {
90+
return err
91+
}
92+
}
93+
94+
var state []byte
95+
if args[1] == "-" {
96+
state, err = io.ReadAll(cmd.InOrStdin())
97+
} else {
98+
state, err = os.ReadFile(args[1])
99+
}
100+
if err != nil {
101+
return err
102+
}
103+
return client.UpdateWorkspaceBuildState(cmd.Context(), build.ID, state)
104+
},
105+
}
106+
cmd.Flags().StringVarP(&buildName, "build", "b", "latest", "Specify a workspace build to target by name.")
107+
return cmd
108+
}
109+
110+
// Resources in a started / stopped state should be generalized anyways.
111+
// SSHing into a stopped instance should provide a

‎cli/state_test.go

Lines changed: 133 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,133 @@
1+
package cli_test
2+
3+
import (
4+
"bytes"
5+
"context"
6+
"os"
7+
"path/filepath"
8+
"testing"
9+
10+
"github.com/stretchr/testify/require"
11+
12+
"github.com/coder/coder/cli/clitest"
13+
"github.com/coder/coder/coderd/coderdtest"
14+
"github.com/coder/coder/provisioner/echo"
15+
"github.com/coder/coder/provisionersdk/proto"
16+
)
17+
18+
func TestStatePull(t *testing.T) {
19+
t.Parallel()
20+
t.Run("File", func(t *testing.T) {
21+
t.Parallel()
22+
client := coderdtest.New(t, nil)
23+
user := coderdtest.CreateFirstUser(t, client)
24+
coderdtest.NewProvisionerDaemon(t, client)
25+
wantState := []byte("some state")
26+
version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
27+
Parse: echo.ParseComplete,
28+
Provision: []*proto.Provision_Response{{
29+
Type: &proto.Provision_Response_Complete{
30+
Complete: &proto.Provision_Complete{
31+
State: wantState,
32+
},
33+
},
34+
}},
35+
})
36+
coderdtest.AwaitTemplateVersionJob(t, client, version.ID)
37+
template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
38+
workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
39+
coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
40+
statefilePath := filepath.Join(t.TempDir(), "state")
41+
cmd, root := clitest.New(t, "state", "pull", workspace.Name, statefilePath)
42+
clitest.SetupConfig(t, client, root)
43+
err := cmd.Execute()
44+
require.NoError(t, err)
45+
gotState, err := os.ReadFile(statefilePath)
46+
require.NoError(t, err)
47+
require.Equal(t, wantState, gotState)
48+
})
49+
t.Run("Stdout", func(t *testing.T) {
50+
t.Parallel()
51+
client := coderdtest.New(t, nil)
52+
user := coderdtest.CreateFirstUser(t, client)
53+
coderdtest.NewProvisionerDaemon(t, client)
54+
wantState := []byte("some state")
55+
version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
56+
Parse: echo.ParseComplete,
57+
Provision: []*proto.Provision_Response{{
58+
Type: &proto.Provision_Response_Complete{
59+
Complete: &proto.Provision_Complete{
60+
State: wantState,
61+
},
62+
},
63+
}},
64+
})
65+
coderdtest.AwaitTemplateVersionJob(t, client, version.ID)
66+
template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
67+
workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
68+
coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
69+
cmd, root := clitest.New(t, "state", "pull", workspace.Name)
70+
var gotState bytes.Buffer
71+
cmd.SetOut(&gotState)
72+
clitest.SetupConfig(t, client, root)
73+
err := cmd.Execute()
74+
require.NoError(t, err)
75+
require.Equal(t, wantState, bytes.TrimSpace(gotState.Bytes()))
76+
})
77+
}
78+
79+
func TestStatePush(t *testing.T) {
80+
t.Parallel()
81+
t.Run("File", func(t *testing.T) {
82+
t.Parallel()
83+
client := coderdtest.New(t, nil)
84+
user := coderdtest.CreateFirstUser(t, client)
85+
coderdtest.NewProvisionerDaemon(t, client)
86+
version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
87+
Parse: echo.ParseComplete,
88+
Provision: echo.ProvisionComplete,
89+
})
90+
coderdtest.AwaitTemplateVersionJob(t, client, version.ID)
91+
template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
92+
workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
93+
coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
94+
stateFile, err := os.CreateTemp(t.TempDir(), "")
95+
require.NoError(t, err)
96+
wantState := []byte("some magic state")
97+
_, err = stateFile.Write(wantState)
98+
require.NoError(t, err)
99+
err = stateFile.Close()
100+
require.NoError(t, err)
101+
cmd, root := clitest.New(t, "state", "push", workspace.Name, stateFile.Name())
102+
clitest.SetupConfig(t, client, root)
103+
err = cmd.Execute()
104+
require.NoError(t, err)
105+
gotState, err := client.WorkspaceBuildState(context.Background(), workspace.LatestBuild.ID)
106+
require.NoError(t, err)
107+
require.Equal(t, wantState, gotState)
108+
})
109+
110+
t.Run("Stdin", func(t *testing.T) {
111+
t.Parallel()
112+
client := coderdtest.New(t, nil)
113+
user := coderdtest.CreateFirstUser(t, client)
114+
coderdtest.NewProvisionerDaemon(t, client)
115+
version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
116+
Parse: echo.ParseComplete,
117+
Provision: echo.ProvisionComplete,
118+
})
119+
coderdtest.AwaitTemplateVersionJob(t, client, version.ID)
120+
template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
121+
workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
122+
coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
123+
wantState := []byte("some magic state")
124+
cmd, root := clitest.New(t, "state", "push", "--build", workspace.LatestBuild.Name, workspace.Name, "-")
125+
clitest.SetupConfig(t, client, root)
126+
cmd.SetIn(bytes.NewReader(wantState))
127+
err := cmd.Execute()
128+
require.NoError(t, err)
129+
gotState, err := client.WorkspaceBuildState(context.Background(), workspace.LatestBuild.ID)
130+
require.NoError(t, err)
131+
require.Equal(t, wantState, gotState)
132+
})
133+
}

‎coderd/coderd.go

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -268,6 +268,8 @@ func New(options *Options) (http.Handler, func()) {
268268
r.Patch("/cancel", api.patchCancelWorkspaceBuild)
269269
r.Get("/logs", api.workspaceBuildLogs)
270270
r.Get("/resources", api.workspaceBuildResources)
271+
r.Get("/state", api.workspaceBuildState)
272+
r.Post("/state", api.postWorkspaceBuildState)
271273
})
272274
})
273275
r.NotFound(site.DefaultHandler().ServeHTTP)

‎coderd/provisionerdaemons.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -95,7 +95,7 @@ func (api *api) provisionerDaemonsListen(rw http.ResponseWriter, r *http.Request
9595
},
9696
})
9797
err = server.Serve(r.Context(), session)
98-
if err != nil {
98+
if err != nil && !xerrors.Is(err, io.EOF) {
9999
api.Logger.Debug(r.Context(), "provisioner daemon disconnected", slog.Error(err))
100100
_ = conn.Close(websocket.StatusInternalError, httpapi.WebsocketCloseSprintf("serve: %s", err))
101101
return

‎coderd/workspacebuilds.go

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package coderd
33
import (
44
"database/sql"
55
"fmt"
6+
"io"
67
"net/http"
78

89
"github.com/coder/coder/coderd/database"
@@ -87,6 +88,42 @@ func (api *api) workspaceBuildLogs(rw http.ResponseWriter, r *http.Request) {
8788
api.provisionerJobLogs(rw, r, job)
8889
}
8990

91+
func (*api) workspaceBuildState(rw http.ResponseWriter, r *http.Request) {
92+
workspaceBuild := httpmw.WorkspaceBuildParam(r)
93+
94+
rw.Header().Set("Content-Type", "application/json")
95+
rw.WriteHeader(http.StatusOK)
96+
_, _ = rw.Write(workspaceBuild.ProvisionerState)
97+
}
98+
99+
func (api *api) postWorkspaceBuildState(rw http.ResponseWriter, r *http.Request) {
100+
workspaceBuild := httpmw.WorkspaceBuildParam(r)
101+
102+
r.Body = http.MaxBytesReader(rw, r.Body, 10*(10<<20))
103+
data, err := io.ReadAll(r.Body)
104+
if err != nil {
105+
httpapi.Write(rw, http.StatusBadRequest, httpapi.Response{
106+
Message: fmt.Sprintf("read state: %s", err),
107+
})
108+
return
109+
}
110+
err = api.Database.UpdateWorkspaceBuildByID(r.Context(), database.UpdateWorkspaceBuildByIDParams{
111+
ID: workspaceBuild.ID,
112+
UpdatedAt: database.Now(),
113+
AfterID: workspaceBuild.AfterID,
114+
ProvisionerState: data,
115+
})
116+
if err != nil {
117+
httpapi.Write(rw, http.StatusInternalServerError, httpapi.Response{
118+
Message: fmt.Sprintf("update workspace build: %s", err),
119+
})
120+
return
121+
}
122+
httpapi.Write(rw, http.StatusOK, httpapi.Response{
123+
Message: "Updated state!",
124+
})
125+
}
126+
90127
func convertWorkspaceBuild(workspaceBuild database.WorkspaceBuild, job codersdk.ProvisionerJob) codersdk.WorkspaceBuild {
91128
//nolint:unconvert
92129
return codersdk.WorkspaceBuild{

‎coderd/workspacebuilds_test.go

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -168,3 +168,61 @@ func TestWorkspaceBuildLogs(t *testing.T) {
168168
}
169169
require.Fail(t, "example message never happened")
170170
}
171+
172+
func TestWorkspaceBuildState(t *testing.T) {
173+
t.Parallel()
174+
client := coderdtest.New(t, nil)
175+
user := coderdtest.CreateFirstUser(t, client)
176+
coderdtest.NewProvisionerDaemon(t, client)
177+
wantState := []byte("some kinda state")
178+
version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
179+
Parse: echo.ParseComplete,
180+
ProvisionDryRun: echo.ProvisionComplete,
181+
Provision: []*proto.Provision_Response{{
182+
Type: &proto.Provision_Response_Complete{
183+
Complete: &proto.Provision_Complete{
184+
State: wantState,
185+
},
186+
},
187+
}},
188+
})
189+
coderdtest.AwaitTemplateVersionJob(t, client, version.ID)
190+
template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
191+
workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
192+
coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
193+
gotState, err := client.WorkspaceBuildState(context.Background(), workspace.LatestBuild.ID)
194+
require.NoError(t, err)
195+
require.Equal(t, wantState, gotState)
196+
}
197+
198+
func TestPostWorkspaceBuildState(t *testing.T) {
199+
t.Parallel()
200+
client := coderdtest.New(t, nil)
201+
user := coderdtest.CreateFirstUser(t, client)
202+
coderdtest.NewProvisionerDaemon(t, client)
203+
wantState := []byte("some kinda state")
204+
version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
205+
Parse: echo.ParseComplete,
206+
ProvisionDryRun: echo.ProvisionComplete,
207+
Provision: []*proto.Provision_Response{{
208+
Type: &proto.Provision_Response_Complete{
209+
Complete: &proto.Provision_Complete{
210+
State: wantState,
211+
},
212+
},
213+
}},
214+
})
215+
coderdtest.AwaitTemplateVersionJob(t, client, version.ID)
216+
template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
217+
workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
218+
coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
219+
gotState, err := client.WorkspaceBuildState(context.Background(), workspace.LatestBuild.ID)
220+
require.NoError(t, err)
221+
require.Equal(t, wantState, gotState)
222+
wantState = []byte("i want a new state")
223+
err = client.UpdateWorkspaceBuildState(context.Background(), workspace.LatestBuild.ID, wantState)
224+
require.NoError(t, err)
225+
gotState, err = client.WorkspaceBuildState(context.Background(), workspace.LatestBuild.ID)
226+
require.NoError(t, err)
227+
require.Equal(t, wantState, gotState)
228+
}

0 commit comments

Comments
 (0)