more explicit cleanup logic and tests for cleaning up the temp dir when loading c1zs (#617)

Author: kans

pkg/dotc1z/c1file.go

@@ -188,7 +188,7 @@ func NewC1ZFile(ctx context.Context, outputFilePath string, opts ...C1ZOption) (
 		opt(options)
 	}
 
-	dbFilePath, err := loadC1z(outputFilePath, options.tmpDir, options.decoderOptions...)
+	dbFilePath, _, err := decompressC1z(outputFilePath, options.tmpDir, options.decoderOptions...)
 	if err != nil {
 		return nil, err
 	}

pkg/dotc1z/dotc1z.go

@@ -50,7 +50,7 @@ func C1ZFileCheckHeader(f io.ReadSeeker) (bool, error) {
 }
 
 func NewExternalC1FileReader(ctx context.Context, tmpDir string, externalResourceC1ZPath string) (connectorstore.Reader, error) {
-	dbFilePath, err := loadC1z(externalResourceC1ZPath, tmpDir)
+	dbFilePath, _, err := decompressC1z(externalResourceC1ZPath, tmpDir)
 	if err != nil {
 		return nil, fmt.Errorf("error loading external resource c1z file: %w", err)
 	}

pkg/dotc1z/file.go

@@ -15,56 +15,83 @@ import (
 	"google.golang.org/grpc/status"
 )
 
-func loadC1z(filePath string, tmpDir string, opts ...DecoderOption) (string, error) {
-	var err error
-	workingDir, err := os.MkdirTemp(tmpDir, "c1z")
+// Note(kans): decompressC1z is unfortunately called to load or create a c1z file so the error handling is rough.
+// It creates its own temporary directory so that it can also do its own cleanup.
+// It returns that directory for verification in tests.
+func decompressC1z(c1zPath string, workingDir string, opts ...DecoderOption) (string, string, error) {
+	tmpDir, err := os.MkdirTemp(workingDir, "c1z")
 	if err != nil {
-		return "", err
+		return "", tmpDir, err
 	}
-	defer func() {
-		if err != nil {
-			if removeErr := os.RemoveAll(workingDir); removeErr != nil {
-				err = errors.Join(err, removeErr)
+
+	var dbFile *os.File
+	var c1zFile *os.File
+	var decoder *decoder
+	cleanupDir := func(e error) error {
+		if decoder != nil {
+			err := decoder.Close()
+			if err != nil {
+				e = errors.Join(e, err)
 			}
 		}
-	}()
-	dbFilePath := filepath.Join(workingDir, "db")
-	dbFile, err := os.Create(dbFilePath)
+		if c1zFile != nil {
+			err := c1zFile.Close()
+			if err != nil {
+				e = errors.Join(e, err)
+			}
+		}
+		if dbFile != nil {
+			err := dbFile.Close()
+			if err != nil {
+				e = errors.Join(e, err)
+			}
+		}
+		if e != nil {
+			err := os.RemoveAll(tmpDir)
+			if err != nil {
+				e = errors.Join(e, err)
+			}
+		}
+		return e
+	}
+
+	dbFilePath := filepath.Join(tmpDir, "db")
+	dbFile, err = os.Create(dbFilePath)
 	if err != nil {
-		return "", err
+		return "", tmpDir, cleanupDir(err)
 	}
-	defer dbFile.Close()
 
-	if stat, err := os.Stat(filePath); err == nil && stat.Size() != 0 {
-		c1zFile, err := os.Open(filePath)
-		if err != nil {
-			return "", err
-		}
-		defer c1zFile.Close()
+	stat, err := os.Stat(c1zPath)
+	if err != nil || stat.Size() == 0 {
+		// TODO(kans): it would be nice to know more about the error....
+		return dbFilePath, tmpDir, cleanupDir(nil)
+	}
 
-		r, err := NewDecoder(c1zFile, opts...)
-		if err != nil {
-			return "", err
-		}
-		_, err = io.Copy(dbFile, r)
-		if err != nil {
-			return "", err
-		}
-		err = r.Close()
-		if err != nil {
-			return "", err
-		}
+	c1zFile, err = os.Open(c1zPath)
+	if err != nil {
+		return "", tmpDir, cleanupDir(err)
+	}
+
+	decoder, err = NewDecoder(c1zFile, opts...)
+	if err != nil {
+		return "", tmpDir, cleanupDir(err)
+	}
+
+	_, err = io.Copy(dbFile, decoder)
+	if err != nil {
+		return "", tmpDir, cleanupDir(err)
 	}
 
 	// CRITICAL: Sync the database file before returning to ensure all
 	// decompressed data is on disk. On filesystems with aggressive caching
 	// (like ZFS with large ARC), SQLite might otherwise open the file and
 	// see incomplete data still in kernel buffers.
-	if err = dbFile.Sync(); err != nil {
-		return "", fmt.Errorf("failed to sync db file: %w", err)
+	err = dbFile.Sync()
+	if err != nil {
+		return "", tmpDir, cleanupDir(err)
 	}
 
-	return dbFilePath, nil
+	return dbFilePath, tmpDir, cleanupDir(nil)
 }
 
 func saveC1z(dbFilePath string, outputFilePath string, encoderConcurrency int) error {

pkg/dotc1z/file_test.go

@@ -12,35 +12,38 @@ import (
 )
 
 func TestLoadC1z(t *testing.T) {
-	tmpDir := t.TempDir()
-
 	t.Run("temp directory cleanup on error", func(t *testing.T) {
 		// Create a file that will cause an error during decoding
+		tmpDir := t.TempDir()
 		invalidFile := filepath.Join(tmpDir, "invalid2.c1z")
 		err := os.WriteFile(invalidFile, []byte("invalid"), 0600)
 		require.NoError(t, err)
 		defer os.Remove(invalidFile)
 
 		// Try to load it - should fail and clean up temp dir
-		dbPath, err := loadC1z(invalidFile, tmpDir)
+		dbPath, workingDir, err := decompressC1z(invalidFile, tmpDir)
 		require.Error(t, err)
 		require.Empty(t, dbPath)
+		// require that the working dir was cleaned up and has no entries
+		require.NoDirExists(t, workingDir)
 	})
 
 	t.Run("custom tmpDir", func(t *testing.T) {
+		tmpDir := t.TempDir()
 		customTmpDir := filepath.Join(tmpDir, "custom")
 		err := os.MkdirAll(customTmpDir, 0755)
 		require.NoError(t, err)
 		defer os.RemoveAll(customTmpDir)
 
 		nonExistentPath := filepath.Join(tmpDir, "nonexistent2.c1z")
-		dbPath, err := loadC1z(nonExistentPath, customTmpDir)
+		dbPath, workingDir, err := decompressC1z(nonExistentPath, customTmpDir)
 		require.NoError(t, err)
 		require.NotEmpty(t, dbPath)
 		require.FileExists(t, dbPath)
 
 		// Verify it was created in the custom tmpDir
 		require.Contains(t, dbPath, customTmpDir)
+		require.Contains(t, workingDir, customTmpDir)
 	})
 }