Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 1aec36d

Browse files
kylecarbskylecarbs
committed
Cleanup code
1 parent 051edda commit 1aec36d

File tree

9 files changed

+188
-68
lines changed

9 files changed

+188
-68
lines changed

‎coderd/coderd.go

+1
Original file line numberDiff line numberDiff line change
@@ -156,6 +156,7 @@ func New(options *Options) (http.Handler, func()) {
156156
r.Get("/", api.provisionerJobResourcesByID)
157157
r.Route("/{workspaceresource}", func(r chi.Router) {
158158
r.Use(httpmw.ExtractWorkspaceResourceParam(options.Database))
159+
r.Get("/", api.provisionerJobResourceByID)
159160
r.Get("/agent", api.workspaceAgentConnectByResource)
160161
})
161162
})

‎coderd/provisionerjobs.go

+42-3
Original file line numberDiff line numberDiff line change
@@ -261,15 +261,54 @@ func (api *api) provisionerJobResourcesByID(rw http.ResponseWriter, r *http.Requ
261261
apiResources = append(apiResources, convertProvisionerJobResource(resource, nil))
262262
continue
263263
}
264-
// TODO: This should be combined.
265-
agents, err := api.Database.GetProvisionerJobAgentsByResourceIDs(r.Context(), []uuid.UUID{resource.ID})
264+
agent, err := api.Database.GetProvisionerJobAgentByResourceID(r.Context(), resource.ID)
265+
if err != nil {
266+
httpapi.Write(rw, http.StatusInternalServerError, httpapi.Response{
267+
Message: fmt.Sprintf("get provisioner job agent: %s", err),
268+
})
269+
return
270+
}
271+
apiAgent, err := convertProvisionerJobAgent(agent)
272+
if err != nil {
273+
httpapi.Write(rw, http.StatusInternalServerError, httpapi.Response{
274+
Message: fmt.Sprintf("convert provisioner job agent: %s", err),
275+
})
276+
return
277+
}
278+
apiResources = append(apiResources, convertProvisionerJobResource(resource, &apiAgent))
279+
}
280+
render.Status(r, http.StatusOK)
281+
render.JSON(rw, r, apiResources)
282+
}
283+
284+
func (api *api) provisionerJobResourceByID(rw http.ResponseWriter, r *http.Request) {
285+
job := httpmw.ProvisionerJobParam(r)
286+
if !convertProvisionerJob(job).Status.Completed() {
287+
httpapi.Write(rw, http.StatusPreconditionFailed, httpapi.Response{
288+
Message: "Job hasn't completed!",
289+
})
290+
return
291+
}
292+
resources, err := api.Database.GetProvisionerJobResourcesByJobID(r.Context(), job.ID)
293+
if err != nil {
294+
httpapi.Write(rw, http.StatusInternalServerError, httpapi.Response{
295+
Message: fmt.Sprintf("get provisioner job resources: %s", err),
296+
})
297+
return
298+
}
299+
apiResources := make([]ProvisionerJobResource, 0)
300+
for _, resource := range resources {
301+
if !resource.AgentID.Valid {
302+
apiResources = append(apiResources, convertProvisionerJobResource(resource, nil))
303+
continue
304+
}
305+
agent, err := api.Database.GetProvisionerJobAgentByResourceID(r.Context(), resource.ID)
266306
if err != nil {
267307
httpapi.Write(rw, http.StatusInternalServerError, httpapi.Response{
268308
Message: fmt.Sprintf("get provisioner job agent: %s", err),
269309
})
270310
return
271311
}
272-
agent := agents[0]
273312
apiAgent, err := convertProvisionerJobAgent(agent)
274313
if err != nil {
275314
httpapi.Write(rw, http.StatusInternalServerError, httpapi.Response{

‎coderd/workspaceagent.go

+35-4
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,16 @@
11
package coderd
22

33
import (
4+
"database/sql"
45
"fmt"
56
"io"
67
"net/http"
8+
"time"
79

8-
"github.com/google/uuid"
910
"github.com/hashicorp/yamux"
1011
"nhooyr.io/websocket"
1112

13+
"github.com/coder/coder/database"
1214
"github.com/coder/coder/httpapi"
1315
"github.com/coder/coder/httpmw"
1416
"github.com/coder/coder/peerbroker"
@@ -27,14 +29,13 @@ func (api *api) workspaceAgentConnectByResource(rw http.ResponseWriter, r *http.
2729
})
2830
return
2931
}
30-
agents, err := api.Database.GetProvisionerJobAgentsByResourceIDs(r.Context(), []uuid.UUID{resource.ID})
32+
agent, err := api.Database.GetProvisionerJobAgentByResourceID(r.Context(), resource.ID)
3133
if err != nil {
3234
httpapi.Write(rw, http.StatusBadRequest, httpapi.Response{
3335
Message: fmt.Sprintf("get provisioner job agent: %s", err),
3436
})
3537
return
3638
}
37-
agent := agents[0]
3839
conn, err := websocket.Accept(rw, r, &websocket.AcceptOptions{
3940
CompressionMode: websocket.CompressionDisabled,
4041
})
@@ -99,5 +100,35 @@ func (api *api) workspaceAgentServe(rw http.ResponseWriter, r *http.Request) {
99100
return
100101
}
101102
defer closer.Close()
102-
<-session.CloseChan()
103+
err = api.Database.UpdateProvisionerJobAgentByID(r.Context(), database.UpdateProvisionerJobAgentByIDParams{
104+
ID: agent.ID,
105+
UpdatedAt: sql.NullTime{
106+
Time: database.Now(),
107+
Valid: true,
108+
},
109+
})
110+
if err != nil {
111+
_ = conn.Close(websocket.StatusAbnormalClosure, err.Error())
112+
return
113+
}
114+
ticker := time.NewTicker(5 * time.Second)
115+
defer ticker.Stop()
116+
for {
117+
select {
118+
case <-session.CloseChan():
119+
return
120+
case <-ticker.C:
121+
err = api.Database.UpdateProvisionerJobAgentByID(r.Context(), database.UpdateProvisionerJobAgentByIDParams{
122+
ID: agent.ID,
123+
UpdatedAt: sql.NullTime{
124+
Time: database.Now(),
125+
Valid: true,
126+
},
127+
})
128+
if err != nil {
129+
_ = conn.Close(websocket.StatusAbnormalClosure, err.Error())
130+
return
131+
}
132+
}
133+
}
103134
}

‎coderd/workspaceagent_test.go

+12-8
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,8 @@ func TestWorkspaceAgentServe(t *testing.T) {
2929
daemonCloser := coderdtest.NewProvisionerDaemon(t, client)
3030
authToken := uuid.NewString()
3131
job := coderdtest.CreateProjectImportJob(t, client, user.Organization, &echo.Responses{
32-
Parse: echo.ParseComplete,
32+
Parse: echo.ParseComplete,
33+
ProvisionDryRun: echo.ProvisionComplete,
3334
Provision: []*proto.Provision_Response{{
3435
Type: &proto.Provision_Response_Complete{
3536
Complete: &proto.Provision_Complete{
@@ -50,26 +51,29 @@ func TestWorkspaceAgentServe(t *testing.T) {
5051
project := coderdtest.CreateProject(t, client, user.Organization, job.ID)
5152
coderdtest.AwaitProjectImportJob(t, client, user.Organization, job.ID)
5253
workspace := coderdtest.CreateWorkspace(t, client, "me", project.ID)
53-
firstHistory, err := client.CreateWorkspaceHistory(context.Background(), "", workspace.Name, coderd.CreateWorkspaceHistoryRequest{
54+
history, err := client.CreateWorkspaceHistory(context.Background(), "", workspace.Name, coderd.CreateWorkspaceHistoryRequest{
5455
ProjectVersionID: project.ActiveVersionID,
5556
Transition: database.WorkspaceTransitionStart,
5657
})
5758
require.NoError(t, err)
58-
coderdtest.AwaitWorkspaceProvisionJob(t, client, user.Organization, firstHistory.ProvisionJobID)
59+
coderdtest.AwaitWorkspaceProvisionJob(t, client, user.Organization, history.ProvisionJobID)
5960
daemonCloser.Close()
60-
resources, err := client.WorkspaceProvisionJobResources(context.Background(), user.Organization, firstHistory.ProvisionJobID)
61-
require.NoError(t, err)
62-
require.Len(t, resources, 1)
6361

6462
agentClient := codersdk.New(client.URL)
6563
agentClient.SessionToken = authToken
6664
agentCloser := agent.New(agentClient.WorkspaceAgentServe, &peer.ConnOptions{
6765
Logger: slogtest.Make(t, nil),
6866
})
6967

70-
time.Sleep(time.Millisecond * 250)
68+
var resources []coderd.ProvisionerJobResource
69+
require.Eventually(t, func() bool {
70+
resources, err = client.WorkspaceProvisionJobResources(context.Background(), user.Organization, history.ProvisionJobID)
71+
require.NoError(t, err)
72+
require.Len(t, resources, 1)
73+
return !resources[0].Agent.UpdatedAt.IsZero()
74+
}, 5*time.Second, 25*time.Millisecond)
7175

72-
workspaceClient, err := client.WorkspaceAgentConnect(context.Background(), user.Organization, firstHistory.ProvisionJobID, resources[0].ID)
76+
workspaceClient, err := client.WorkspaceAgentConnect(context.Background(), user.Organization, history.ProvisionJobID, resources[0].ID)
7377
require.NoError(t, err)
7478
stream, err := workspaceClient.NegotiateConnection(context.Background())
7579
require.NoError(t, err)

‎database/databasefake/databasefake.go

+20-12
Original file line numberDiff line numberDiff line change
@@ -512,8 +512,7 @@ func (q *fakeQuerier) GetProvisionerJobAgentByAuthToken(_ context.Context, authT
512512
q.mutex.Lock()
513513
defer q.mutex.Unlock()
514514

515-
for i := len(q.provisionerJobAgent) - 1; i >= 0; i-- {
516-
agent := q.provisionerJobAgent[i]
515+
for _, agent := range q.provisionerJobAgent {
517516
if agent.AuthToken.String() == authToken.String() {
518517
return agent, nil
519518
}
@@ -535,22 +534,16 @@ func (q *fakeQuerier) GetProvisionerJobAgentByInstanceID(_ context.Context, inst
535534
return database.ProvisionerJobAgent{}, sql.ErrNoRows
536535
}
537536

538-
func (q *fakeQuerier) GetProvisionerJobAgentsByResourceIDs(_ context.Context, ids []uuid.UUID) ([]database.ProvisionerJobAgent, error) {
537+
func (q *fakeQuerier) GetProvisionerJobAgentByResourceID(ctx context.Context, resourceID uuid.UUID) (database.ProvisionerJobAgent, error) {
539538
q.mutex.Lock()
540539
defer q.mutex.Unlock()
541540

542-
agents := make([]database.ProvisionerJobAgent, 0)
543541
for _, agent := range q.provisionerJobAgent {
544-
for _, id := range ids {
545-
if agent.ResourceID.String() == id.String() {
546-
agents = append(agents, agent)
547-
}
542+
if agent.ResourceID.String() == resourceID.String() {
543+
return agent, nil
548544
}
549545
}
550-
if len(agents) == 0 {
551-
return nil, sql.ErrNoRows
552-
}
553-
return agents, nil
546+
return database.ProvisionerJobAgent{}, sql.ErrNoRows
554547
}
555548

556549
func (q *fakeQuerier) GetProvisionerDaemonByID(_ context.Context, id uuid.UUID) (database.ProvisionerDaemon, error) {
@@ -969,6 +962,21 @@ func (q *fakeQuerier) UpdateProvisionerDaemonByID(_ context.Context, arg databas
969962
return sql.ErrNoRows
970963
}
971964

965+
func (q *fakeQuerier) UpdateProvisionerJobAgentByID(ctx context.Context, arg database.UpdateProvisionerJobAgentByIDParams) error {
966+
q.mutex.Lock()
967+
defer q.mutex.Unlock()
968+
969+
for index, agent := range q.provisionerJobAgent {
970+
if agent.ID.String() != arg.ID.String() {
971+
continue
972+
}
973+
agent.UpdatedAt = arg.UpdatedAt
974+
q.provisionerJobAgent[index] = agent
975+
return nil
976+
}
977+
return sql.ErrNoRows
978+
}
979+
972980
func (q *fakeQuerier) UpdateProvisionerJobByID(_ context.Context, arg database.UpdateProvisionerJobByIDParams) error {
973981
q.mutex.Lock()
974982
defer q.mutex.Unlock()

‎database/querier.go

+2-1
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎database/query.sql

+10-2
Original file line numberDiff line numberDiff line change
@@ -354,13 +354,13 @@ FROM
354354
WHERE
355355
job_id = $1;
356356

357-
-- name: GetProvisionerJobAgentsByResourceIDs :many
357+
-- name: GetProvisionerJobAgentByResourceID :one
358358
SELECT
359359
*
360360
FROM
361361
provisioner_job_agent
362362
WHERE
363-
resource_id = ANY(@ids :: uuid [ ]);
363+
resource_id = $1;
364364

365365
-- name: InsertAPIKey :one
366366
INSERT INTO
@@ -660,6 +660,14 @@ SET
660660
WHERE
661661
id = $1;
662662

663+
-- name: UpdateProvisionerJobAgentByID :exec
664+
UPDATE
665+
provisioner_job_agent
666+
SET
667+
updated_at = $2
668+
WHERE
669+
id = $1;
670+
663671
-- name: UpdateWorkspaceHistoryByID :exec
664672
UPDATE
665673
workspace_history

‎database/query.sql.go

+37-34
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)