fix: fetch provisioner logs after end of logs message

kylecarbs · kylecarbs · commit c64dca2ed826 · 2023-03-07T20:57:43.000Z
I think this should _actually_ fix it. See https://github.com/coder/coder/actions/runs/4358242625/jobs/7618562167 The problem was that this loop is ran async so even though messages are sent in order, we were processing too slowly and the end of logs was published first.
diff --git a/coderd/provisionerjobs.go b/coderd/provisionerjobs.go
@@ -13,6 +13,7 @@ import (
 	"time"
 
 	"github.com/google/uuid"
+	"go.uber.org/atomic"
 	"nhooyr.io/websocket"
 
 	"cdr.dev/slog"
@@ -374,26 +375,42 @@ func (api *API) followLogs(actor rbac.Subject, jobID uuid.UUID) (<-chan *databas
 	logger := api.Logger.With(slog.F("job_id", jobID))
 
 	var (
-		closed       = make(chan struct{})
-		bufferedLogs = make(chan *database.ProvisionerJobLog, 128)
-		logMut       = &sync.Mutex{}
+		closeMutex    sync.Mutex
+		bufferedLogs  = make(chan *database.ProvisionerJobLog, 128)
+		endOfLogs     atomic.Bool
+		lastSentLogID atomic.Int64
 	)
+
+	sendLog := func(log *database.ProvisionerJobLog) {
+		closeMutex.Lock()
+		defer closeMutex.Unlock()
+		select {
+		case bufferedLogs <- log:
+			logger.Debug(context.Background(), "subscribe buffered log", slog.F("stage", log.Stage))
+			lastSentLogID.Store(log.ID)
+		default:
+			// If this overflows users could miss logs streaming. This can happen
+			// we get a lot of logs and consumer isn't keeping up.  We don't want to block the pubsub,
+			// so just drop them.
+			logger.Warn(context.Background(), "provisioner job log overflowing channel")
+		}
+	}
+
 	closeSubscribe, err := api.Pubsub.Subscribe(
 		provisionerJobLogsChannel(jobID),
 		func(ctx context.Context, message []byte) {
-			select {
-			case <-closed:
+			if endOfLogs.Load() {
+				// Logs have already ended!
 				return
-			default:
 			}
-
 			jlMsg := provisionerJobLogsMessage{}
 			err := json.Unmarshal(message, &jlMsg)
 			if err != nil {
 				logger.Warn(ctx, "invalid provisioner job log on channel", slog.Error(err))
 				return
 			}
 
+			// CreatedAfter is sent when logs are streaming!
 			if jlMsg.CreatedAfter != 0 {
 				logs, err := api.Database.GetProvisionerLogsByIDBetween(dbauthz.As(ctx, actor), database.GetProvisionerLogsByIDBetweenParams{
 					JobID:        jobID,
@@ -403,54 +420,50 @@ func (api *API) followLogs(actor rbac.Subject, jobID uuid.UUID) (<-chan *databas
 					logger.Warn(ctx, "get provisioner logs", slog.Error(err))
 					return
 				}
-
 				for _, log := range logs {
-					// Sadly we have to use a mutex here because events may be
-					// handled out of order due to golang goroutine scheduling
-					// semantics (even though Postgres guarantees ordering of
-					// notifications).
-					logMut.Lock()
-					select {
-					case <-closed:
-						logMut.Unlock()
+					if endOfLogs.Load() {
+						// An end of logs message came in while we were fetching
+						// logs or processing them!
 						return
-					default:
 					}
 					log := log
-					select {
-					case bufferedLogs <- &log:
-						logger.Debug(ctx, "subscribe buffered log", slog.F("stage", log.Stage))
-					default:
-						// If this overflows users could miss logs streaming. This can happen
-						// we get a lot of logs and consumer isn't keeping up.  We don't want to block the pubsub,
-						// so just drop them.
-						logger.Warn(ctx, "provisioner job log overflowing channel")
-					}
-					logMut.Unlock()
+					sendLog(&log)
 				}
 			}
 
+			// EndOfLogs is sent when logs are done streaming.
+			// We don't want to end the stream until we've sent all the logs,
+			// so we fetch logs after the last ID we've seen and send them!
 			if jlMsg.EndOfLogs {
-				// This mutex is to guard double-closes.
-				logMut.Lock()
-				select {
-				case <-closed:
-					logMut.Unlock()
+				endOfLogs.Store(true)
+				logs, err := api.Database.GetProvisionerLogsByIDBetween(dbauthz.As(ctx, actor), database.GetProvisionerLogsByIDBetweenParams{
+					JobID:        jobID,
+					CreatedAfter: lastSentLogID.Load(),
+				})
+				if err != nil {
+					logger.Warn(ctx, "get provisioner logs", slog.Error(err))
 					return
-				default:
+				}
+				for _, log := range logs {
+					log := log
+					sendLog(&log)
 				}
 				logger.Debug(ctx, "got End of Logs")
+				closeMutex.Lock()
 				bufferedLogs <- nil
-				logMut.Unlock()
+				closeMutex.Unlock()
 			}
+
+			lastSentLogID.Store(jlMsg.CreatedAfter)
 		},
 	)
 	if err != nil {
 		return nil, nil, err
 	}
 	return bufferedLogs, func() {
+		closeMutex.Lock()
+		defer closeMutex.Unlock()
 		closeSubscribe()
-		close(closed)
 		close(bufferedLogs)
 	}, nil
 }
