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

* test: fix cleanup order on provisioner daemon work dir

* Reduce the test race condition

coderd/coderdtest/coderdtest.go

@@ -459,6 +459,12 @@ 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() {
@@ -466,8 +472,7 @@ func NewProvisionerDaemon(t testing.TB, coderAPI *coderd.API) io.Closer {
 		_ = 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,

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)