chore: fix buffered provisioner job logs close flake

kylecarbs · kylecarbs · commit 9185c0693d5e · 2023-03-07T19:34:03.000Z
See https://github.com/coder/coder/actions/runs/4357599919/jobs/7617111287
diff --git a/coderd/provisionerjobs.go b/coderd/provisionerjobs.go
@@ -50,8 +50,9 @@ func (api *API) provisionerJobLogs(rw http.ResponseWriter, r *http.Request, job
 	// between the end of our query and the start of the subscription.  We might get duplicates, so we'll keep track
 	// of processed IDs.
 	var bufferedLogs <-chan database.ProvisionerJobLog
+	var bufferedLogClosed chan struct{}
 	if follow {
-		bl, closeFollow, err := api.followLogs(actor, job.ID)
+		bl, closed, closeFollow, err := api.followLogs(actor, job.ID)
 		if err != nil {
 			httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
 				Message: "Internal error watching provisioner logs.",
@@ -61,6 +62,7 @@ func (api *API) provisionerJobLogs(rw http.ResponseWriter, r *http.Request, job
 		}
 		defer closeFollow()
 		bufferedLogs = bl
+		bufferedLogClosed = closed
 
 		// Next query the job itself to see if it is complete.  If so, the historical query to the database will return
 		// the full set of logs.  It's a little sad to have to query the job again, given that our caller definitely
@@ -188,6 +190,12 @@ func (api *API) provisionerJobLogs(rw http.ResponseWriter, r *http.Request, job
 					return
 				}
 			}
+		// This must come after the buffered logs check above, because
+		// it's important that all buffered logs are flushed before we
+		// close the connection!
+		case <-bufferedLogClosed:
+			logger.Debug(context.Background(), "done with buffered logs")
+			return
 		}
 	}
 }
@@ -369,7 +377,7 @@ type provisionerJobLogsMessage struct {
 	EndOfLogs    bool  `json:"end_of_logs,omitempty"`
 }
 
-func (api *API) followLogs(actor rbac.Subject, jobID uuid.UUID) (<-chan database.ProvisionerJobLog, func(), error) {
+func (api *API) followLogs(actor rbac.Subject, jobID uuid.UUID) (<-chan database.ProvisionerJobLog, chan struct{}, func(), error) {
 	logger := api.Logger.With(slog.F("job_id", jobID))
 
 	var (
@@ -447,7 +455,7 @@ func (api *API) followLogs(actor rbac.Subject, jobID uuid.UUID) (<-chan database
 		},
 	)
 	if err != nil {
-		return nil, nil, err
+		return nil, nil, nil, err
 	}
-	return bufferedLogs, closeSubscribe, nil
+	return bufferedLogs, closed, closeSubscribe, nil
 }

Original file line number	Diff line number	Diff line change
`@@ -50,8 +50,9 @@ func (api API) provisionerJobLogs(rw http.ResponseWriter, r http.Request, job`
`50`	`50`	`// between the end of our query and the start of the subscription. We might get duplicates, so we'll keep track`
`51`	`51`	`// of processed IDs.`
`52`	`52`	`var bufferedLogs <-chan database.ProvisionerJobLog`
	`53`	`+ var bufferedLogClosed chan struct{}`
`53`	`54`	`if follow {`
`54`		`- bl, closeFollow, err := api.followLogs(actor, job.ID)`
	`55`	`+ bl, closed, closeFollow, err := api.followLogs(actor, job.ID)`
`55`	`56`	`if err != nil {`
`56`	`57`	`httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{`
`57`	`58`	`Message: "Internal error watching provisioner logs.",`
`@@ -61,6 +62,7 @@ func (api API) provisionerJobLogs(rw http.ResponseWriter, r http.Request, job`
`61`	`62`	`}`
`62`	`63`	`defer closeFollow()`
`63`	`64`	`bufferedLogs = bl`
	`65`	`+ bufferedLogClosed = closed`
`64`	`66`
`65`	`67`	`// Next query the job itself to see if it is complete. If so, the historical query to the database will return`
`66`	`68`	`// the full set of logs. It's a little sad to have to query the job again, given that our caller definitely`
`@@ -188,6 +190,12 @@ func (api API) provisionerJobLogs(rw http.ResponseWriter, r http.Request, job`
`188`	`190`	`return`
`189`	`191`	`}`
`190`	`192`	`}`
	`193`	`+ // This must come after the buffered logs check above, because`
	`194`	`+ // it's important that all buffered logs are flushed before we`
	`195`	`+ // close the connection!`
	`196`	`+ case <-bufferedLogClosed:`
	`197`	`+ logger.Debug(context.Background(), "done with buffered logs")`
	`198`	`+ return`
`191`	`199`	`}`
`192`	`200`	`}`
`193`	`201`	`}`
`@@ -369,7 +377,7 @@ type provisionerJobLogsMessage struct {`
`369`	`377`	EndOfLogs bool `json:"end_of_logs,omitempty"`
`370`	`378`	`}`
`371`	`379`
`372`		`-func (api *API) followLogs(actor rbac.Subject, jobID uuid.UUID) (<-chan database.ProvisionerJobLog, func(), error) {`
	`380`	`+func (api *API) followLogs(actor rbac.Subject, jobID uuid.UUID) (<-chan database.ProvisionerJobLog, chan struct{}, func(), error) {`
`373`	`381`	`logger := api.Logger.With(slog.F("job_id", jobID))`
`374`	`382`
`375`	`383`	`var (`
`@@ -447,7 +455,7 @@ func (api *API) followLogs(actor rbac.Subject, jobID uuid.UUID) (<-chan database`
`447`	`455`	`},`
`448`	`456`	`)`
`449`	`457`	`if err != nil {`
`450`		`- return nil, nil, err`
	`458`	`+ return nil, nil, nil, err`
`451`	`459`	`}`
`452`		`- return bufferedLogs, closeSubscribe, nil`
	`460`	`+ return bufferedLogs, closed, closeSubscribe, nil`
`453`	`461`	`}`