Agate/performance_test.go

package agate

import (
	"context"
	"fmt"
	"os"
	"path/filepath"
	"testing"
	"time"
)

// BenchmarkCreateSnapshot benchmarks the performance of creating snapshots with different numbers of files
func BenchmarkCreateSnapshot(b *testing.B) {
	// Skip in short mode
	if testing.Short() {
		b.Skip("Skipping benchmark in short mode")
	}

	// Test with different numbers of files
	fileCounts := []int{10, 100, 1000}
	for _, fileCount := range fileCounts {
		b.Run(fmt.Sprintf("Files-%d", fileCount), func(b *testing.B) {
			// Create a temporary directory for tests
			tempDir, err := os.MkdirTemp("", "agate-bench-*")
			if err != nil {
				b.Fatalf("Failed to create temp directory: %v", err)
			}
			defer os.RemoveAll(tempDir)

			// Create a data directory
			dataDir := filepath.Join(tempDir, "data")
			if err := os.MkdirAll(dataDir, 0755); err != nil {
				b.Fatalf("Failed to create data directory: %v", err)
			}

			// Create test files
			createBenchmarkFiles(b, dataDir, fileCount, 1024) // 1 KB per file

			// Create Agate options
			options := AgateOptions{
				WorkDir: dataDir,
				OpenFunc: func(dir string) error {
					return nil
				},
				CloseFunc: func() error {
					return nil
				},
			}

			// Create Agate instance
			ag, err := New(options)
			if err != nil {
				b.Fatalf("Failed to create Agate instance: %v", err)
			}
			defer ag.Close()

			// Reset the timer before the benchmark loop
			b.ResetTimer()

			// Run the benchmark
			for i := 0; i < b.N; i++ {
				ctx := context.Background()
				_, err := ag.SaveSnapshot(ctx, fmt.Sprintf("Benchmark Snapshot %d", i), "")
				if err != nil {
					b.Fatalf("Failed to create snapshot: %v", err)
				}
			}
		})
	}
}

// BenchmarkRestoreSnapshot benchmarks the performance of restoring snapshots with different numbers of files
func BenchmarkRestoreSnapshot(b *testing.B) {
	// Skip in short mode
	if testing.Short() {
		b.Skip("Skipping benchmark in short mode")
	}

	// Test with different numbers of files
	fileCounts := []int{10, 100, 1000}
	for _, fileCount := range fileCounts {
		b.Run(fmt.Sprintf("Files-%d", fileCount), func(b *testing.B) {
			// Create a temporary directory for tests
			tempDir, err := os.MkdirTemp("", "agate-bench-*")
			if err != nil {
				b.Fatalf("Failed to create temp directory: %v", err)
			}
			defer os.RemoveAll(tempDir)

			// Create a data directory
			dataDir := filepath.Join(tempDir, "data")
			if err := os.MkdirAll(dataDir, 0755); err != nil {
				b.Fatalf("Failed to create data directory: %v", err)
			}

			// Create test files
			createBenchmarkFiles(b, dataDir, fileCount, 1024) // 1 KB per file

			// Create Agate options
			options := AgateOptions{
				WorkDir: dataDir,
				OpenFunc: func(dir string) error {
					return nil
				},
				CloseFunc: func() error {
					return nil
				},
			}

			// Create Agate instance
			ag, err := New(options)
			if err != nil {
				b.Fatalf("Failed to create Agate instance: %v", err)
			}
			defer ag.Close()

			// Create a snapshot
			ctx := context.Background()
			snapshotID, err := ag.SaveSnapshot(ctx, "Benchmark Snapshot", "")
			if err != nil {
				b.Fatalf("Failed to create snapshot: %v", err)
			}

			// Modify some files
			for i := 0; i < fileCount/2; i++ {
				filePath := filepath.Join(dataDir, fmt.Sprintf("file_%d.txt", i))
				if err := os.WriteFile(filePath, []byte(fmt.Sprintf("Modified content %d", i)), 0644); err != nil {
					b.Fatalf("Failed to modify file: %v", err)
				}
			}

			// Reset the timer before the benchmark loop
			b.ResetTimer()

			// Run the benchmark
			for i := 0; i < b.N; i++ {
				err := ag.RestoreSnapshot(ctx, snapshotID)
				if err != nil {
					b.Fatalf("Failed to restore snapshot: %v", err)
				}
			}
		})
	}
}

// BenchmarkLargeFiles benchmarks the performance of creating and restoring snapshots with large files
func BenchmarkLargeFiles(b *testing.B) {
	// Skip in short mode
	if testing.Short() {
		b.Skip("Skipping benchmark in short mode")
	}

	// Test with different file sizes
	fileSizes := []int{1 * 1024 * 1024, 10 * 1024 * 1024, 100 * 1024 * 1024} // 1 MB, 10 MB, 100 MB
	for _, fileSize := range fileSizes {
		b.Run(fmt.Sprintf("Size-%dMB", fileSize/(1024*1024)), func(b *testing.B) {
			// Create a temporary directory for tests
			tempDir, err := os.MkdirTemp("", "agate-bench-*")
			if err != nil {
				b.Fatalf("Failed to create temp directory: %v", err)
			}
			defer os.RemoveAll(tempDir)

			// Create a data directory
			dataDir := filepath.Join(tempDir, "data")
			if err := os.MkdirAll(dataDir, 0755); err != nil {
				b.Fatalf("Failed to create data directory: %v", err)
			}

			// Create a large file
			largeFilePath := filepath.Join(dataDir, "large_file.bin")
			createLargeFile(b, largeFilePath, fileSize)

			// Create Agate options
			options := AgateOptions{
				WorkDir: dataDir,
				OpenFunc: func(dir string) error {
					return nil
				},
				CloseFunc: func() error {
					return nil
				},
			}

			// Create Agate instance
			ag, err := New(options)
			if err != nil {
				b.Fatalf("Failed to create Agate instance: %v", err)
			}
			defer ag.Close()

			// Create a snapshot
			ctx := context.Background()

			// Measure snapshot creation time
			b.Run("Create", func(b *testing.B) {
				b.ResetTimer()
				for i := 0; i < b.N; i++ {
					_, err := ag.SaveSnapshot(ctx, fmt.Sprintf("Large File Snapshot %d", i), "")
					if err != nil {
						b.Fatalf("Failed to create snapshot: %v", err)
					}
				}
			})

			// Create a snapshot for restoration benchmark
			snapshotID, err := ag.SaveSnapshot(ctx, "Large File Snapshot", "")
			if err != nil {
				b.Fatalf("Failed to create snapshot: %v", err)
			}

			// Modify the large file
			if err := os.WriteFile(largeFilePath, []byte("Modified content"), 0644); err != nil {
				b.Fatalf("Failed to modify large file: %v", err)
			}

			// Measure snapshot restoration time
			b.Run("Restore", func(b *testing.B) {
				b.ResetTimer()
				for i := 0; i < b.N; i++ {
					err := ag.RestoreSnapshot(ctx, snapshotID)
					if err != nil {
						b.Fatalf("Failed to restore snapshot: %v", err)
					}
				}
			})
		})
	}
}

// TestPerformanceMetrics runs performance tests and reports metrics
func TestPerformanceMetrics(t *testing.T) {
	// Skip in short mode
	if testing.Short() {
		t.Skip("Skipping performance metrics test in short mode")
	}

	// Create a temporary directory for tests
	tempDir, err := os.MkdirTemp("", "agate-perf-*")
	if err != nil {
		t.Fatalf("Failed to create temp directory: %v", err)
	}
	defer os.RemoveAll(tempDir)

	// Create a data directory
	dataDir := filepath.Join(tempDir, "data")
	if err := os.MkdirAll(dataDir, 0755); err != nil {
		t.Fatalf("Failed to create data directory: %v", err)
	}

	// Test with different numbers of files
	fileCounts := []int{10, 100, 1000}
	for _, fileCount := range fileCounts {
		t.Run(fmt.Sprintf("Files-%d", fileCount), func(t *testing.T) {
			// Create test files
			createBenchmarkFiles(t, dataDir, fileCount, 1024) // 1 KB per file

			// Create Agate options
			options := AgateOptions{
				WorkDir: dataDir,
				OpenFunc: func(dir string) error {
					return nil
				},
				CloseFunc: func() error {
					return nil
				},
			}

			// Create Agate instance
			ag, err := New(options)
			if err != nil {
				t.Fatalf("Failed to create Agate instance: %v", err)
			}
			defer ag.Close()

			// Measure snapshot creation time
			ctx := context.Background()
			startTime := time.Now()
			snapshotID, err := ag.SaveSnapshot(ctx, "Performance Test Snapshot", "")
			if err != nil {
				t.Fatalf("Failed to create snapshot: %v", err)
			}
			createDuration := time.Since(startTime)
			t.Logf("Created snapshot with %d files in %v (%.2f files/sec)", fileCount, createDuration, float64(fileCount)/createDuration.Seconds())

			// Modify some files
			for i := 0; i < fileCount/2; i++ {
				filePath := filepath.Join(dataDir, fmt.Sprintf("file_%d.txt", i))
				if err := os.WriteFile(filePath, []byte(fmt.Sprintf("Modified content %d", i)), 0644); err != nil {
					t.Fatalf("Failed to modify file: %v", err)
				}
			}

			// Measure snapshot restoration time
			startTime = time.Now()
			err = ag.RestoreSnapshot(ctx, snapshotID)
			if err != nil {
				t.Fatalf("Failed to restore snapshot: %v", err)
			}
			restoreDuration := time.Since(startTime)
			t.Logf("Restored snapshot with %d files in %v (%.2f files/sec)", fileCount, restoreDuration, float64(fileCount)/restoreDuration.Seconds())
		})
	}
}

// Helper function to create benchmark files
func createBenchmarkFiles(tb testing.TB, dir string, count, size int) {
	tb.Helper()

	// Create files with sequential names
	for i := 0; i < count; i++ {
		filePath := filepath.Join(dir, fmt.Sprintf("file_%d.txt", i))

		// Create content of specified size
		content := make([]byte, size)
		for j := 0; j < size; j++ {
			content[j] = byte(j % 256)
		}

		if err := os.WriteFile(filePath, content, 0644); err != nil {
			tb.Fatalf("Failed to create benchmark file %s: %v", filePath, err)
		}
	}
}

// Helper function to create a large file
func createLargeFile(tb testing.TB, path string, size int) {
	tb.Helper()

	// Create the file
	file, err := os.Create(path)
	if err != nil {
		tb.Fatalf("Failed to create large file: %v", err)
	}
	defer file.Close()

	// Create a buffer with a pattern
	bufferSize := 8192 // 8 KB buffer
	buffer := make([]byte, bufferSize)
	for i := 0; i < bufferSize; i++ {
		buffer[i] = byte(i % 256)
	}

	// Write the buffer multiple times to reach the desired size
	bytesWritten := 0
	for bytesWritten < size {
		n, err := file.Write(buffer)
		if err != nil {
			tb.Fatalf("Failed to write to large file: %v", err)
		}
		bytesWritten += n
	}
}