test: fix cleanup order on provisioner daemon work dir (#9668)

Emyrk · web-flow · commit 0e4d6896e30d · 2023-09-13T13:13:08.000-05:00
* test: fix cleanup order on provisioner daemon work dir

* Reduce the test race condition
diff --git a/coderd/coderdtest/coderdtest.go b/coderd/coderdtest/coderdtest.go
@@ -459,15 +459,20 @@ func NewWithAPI(t testing.TB, options *Options) (*codersdk.Client, io.Closer, *c
 func NewProvisionerDaemon(t testing.TB, coderAPI *coderd.API) io.Closer {
 	t.Helper()
 
+	// t.Cleanup runs in last added, first called order. t.TempDir() will delete
+	// the directory on cleanup, so we want to make sure the echoServer is closed
+	// before we go ahead an attempt to delete it's work directory.
+	// seems t.TempDir() is not safe to call from a different goroutine
+	workDir := t.TempDir()
+
 	echoClient, echoServer := provisionersdk.MemTransportPipe()
 	ctx, cancelFunc := context.WithCancel(context.Background())
 	t.Cleanup(func() {
 		_ = echoClient.Close()
 		_ = echoServer.Close()
 		cancelFunc()
 	})
-	// seems t.TempDir() is not safe to call from a different goroutine
-	workDir := t.TempDir()
+
 	go func() {
 		err := echo.Serve(ctx, &provisionersdk.ServeOptions{
 			Listener:      echoServer,
diff --git a/provisionersdk/session.go b/provisionersdk/session.go
@@ -230,6 +230,15 @@ func (s *Session) extractArchive() error {
 		if mode == 0 {
 			mode = 0o600
 		}
+
+		// Always check for context cancellation before reading the next header.
+		// This is mainly important for unit tests, since a canceled context means
+		// the underlying directory is going to be deleted. There still exists
+		// the small race condition that the context is cancelled after this, and
+		// before the disk write.
+		if ctx.Err() != nil {
+			return xerrors.Errorf("context canceled: %w", ctx.Err())
+		}
 		switch header.Typeflag {
 		case tar.TypeDir:
 			err = os.MkdirAll(headerPath, mode)
