276 lines
7.6 KiB
Go
276 lines
7.6 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"log"
|
|
"os"
|
|
"path/filepath"
|
|
"runtime"
|
|
"sync"
|
|
"sync/atomic"
|
|
"time"
|
|
)
|
|
|
|
// RemoveStats tracks statistics about the file removal process
|
|
type RemoveStats struct {
|
|
Total int64 // Total number of files/directories processed
|
|
Success int64 // Number of successfully removed files/directories
|
|
Failed int64 // Number of files/directories that failed to be removed
|
|
Errors []error // Collection of errors encountered during removal
|
|
mu sync.Mutex // Mutex to protect concurrent access to Errors slice
|
|
}
|
|
|
|
// RemoveOptions configures the file removal process
|
|
type RemoveOptions struct {
|
|
Workers int // Number of concurrent workers (defaults to CPU count)
|
|
Logger func(format string, args ...any) // Logger function
|
|
Retries int // Number of retries for failed removals
|
|
Stats *RemoveStats // Statistics collector (optional)
|
|
Progress func(current int64) // Progress callback function (optional)
|
|
}
|
|
|
|
// removeResult represents the result of a single file/directory removal operation
|
|
type removeResult struct {
|
|
Path string // Path that was processed
|
|
Err error // Error if removal failed, nil if successful
|
|
}
|
|
|
|
// Remove recursively removes a directory and all its contents with concurrency support
|
|
func Remove(ctx context.Context, path string, opts RemoveOptions) error {
|
|
// Apply default options if not specified
|
|
if opts.Workers <= 0 {
|
|
opts.Workers = runtime.NumCPU()
|
|
}
|
|
if opts.Logger == nil {
|
|
opts.Logger = log.Printf
|
|
}
|
|
if opts.Retries < 0 {
|
|
opts.Retries = 0
|
|
}
|
|
if opts.Stats == nil {
|
|
opts.Stats = &RemoveStats{}
|
|
}
|
|
|
|
// Validate path
|
|
if path == "" {
|
|
return fmt.Errorf("empty path provided")
|
|
}
|
|
|
|
// Safety check: prevent removal of critical directories
|
|
absPath, err := filepath.Abs(path)
|
|
if err == nil {
|
|
// List of critical paths that should never be removed
|
|
criticalPaths := []string{
|
|
"/", "/bin", "/boot", "/dev", "/etc", "/home", "/lib", "/lib64",
|
|
"/media", "/mnt", "/opt", "/proc", "/root", "/run", "/sbin",
|
|
"/srv", "/sys", "/tmp", "/usr", "/var",
|
|
"C:\\", "C:\\Windows", "C:\\Program Files", "C:\\Program Files (x86)",
|
|
"C:\\Users",
|
|
}
|
|
|
|
for _, critical := range criticalPaths {
|
|
if absPath == critical {
|
|
return fmt.Errorf("refusing to remove critical system directory: %s", absPath)
|
|
}
|
|
}
|
|
}
|
|
|
|
// Check if path exists
|
|
info, err := os.Stat(path)
|
|
if os.IsNotExist(err) {
|
|
opts.Logger("[CLEAN] Path does not exist, skipping: %s", path)
|
|
return nil
|
|
}
|
|
if err != nil {
|
|
return fmt.Errorf("failed to access path %s: %w", path, err)
|
|
}
|
|
|
|
// Check if it's a directory
|
|
if !info.IsDir() {
|
|
opts.Logger("[CLEAN] Path is not a directory, removing file: %s", path)
|
|
if err := os.Remove(path); err != nil {
|
|
return fmt.Errorf("failed to remove file %s: %w", path, err)
|
|
}
|
|
if opts.Stats != nil {
|
|
atomic.AddInt64(&opts.Stats.Total, 1)
|
|
atomic.AddInt64(&opts.Stats.Success, 1)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
opts.Logger("[CLEAN] Starting removal of directory: %s with %d workers", path, opts.Workers)
|
|
|
|
// Create channels for work distribution and result collection
|
|
workChan := make(chan string, opts.Workers*2)
|
|
resultChan := make(chan removeResult, opts.Workers*2)
|
|
|
|
// Start worker goroutines
|
|
var wg sync.WaitGroup
|
|
for i := 0; i < opts.Workers; i++ {
|
|
wg.Add(1)
|
|
go removeWorker(ctx, &wg, workChan, resultChan, opts)
|
|
}
|
|
|
|
// Start directory traversal
|
|
go func() {
|
|
defer close(workChan)
|
|
if err := walkDir(ctx, path, workChan, opts); err != nil {
|
|
opts.Logger("[CLEAN] Error walking directory %s: %v", path, err)
|
|
}
|
|
}()
|
|
|
|
// Wait for workers to finish and close result channel
|
|
go func() {
|
|
wg.Wait()
|
|
close(resultChan)
|
|
}()
|
|
|
|
// Process results
|
|
var firstErr error
|
|
for res := range resultChan {
|
|
atomic.AddInt64(&opts.Stats.Total, 1)
|
|
if res.Err != nil {
|
|
atomic.AddInt64(&opts.Stats.Failed, 1)
|
|
opts.Stats.mu.Lock()
|
|
opts.Stats.Errors = append(opts.Stats.Errors, res.Err)
|
|
opts.Stats.mu.Unlock()
|
|
if firstErr == nil {
|
|
firstErr = res.Err
|
|
}
|
|
} else {
|
|
atomic.AddInt64(&opts.Stats.Success, 1)
|
|
}
|
|
if opts.Progress != nil {
|
|
opts.Progress(atomic.LoadInt64(&opts.Stats.Total))
|
|
}
|
|
}
|
|
|
|
// Try to remove the root directory itself
|
|
if err := tryRemoveRoot(ctx, path, opts); err != nil && firstErr == nil {
|
|
firstErr = err
|
|
atomic.AddInt64(&opts.Stats.Failed, 1)
|
|
opts.Stats.mu.Lock()
|
|
opts.Stats.Errors = append(opts.Stats.Errors, err)
|
|
opts.Stats.mu.Unlock()
|
|
}
|
|
|
|
opts.Logger("[CLEAN] Completed removal of %s: Total=%d Success=%d Failed=%d",
|
|
path, opts.Stats.Total, opts.Stats.Success, opts.Stats.Failed)
|
|
|
|
return firstErr
|
|
}
|
|
|
|
// walkDir traverses the directory tree and sends paths to the work channel
|
|
// Uses depth-first, reverse order to ensure children are processed before parents
|
|
func walkDir(ctx context.Context, root string, workChan chan<- string, opts RemoveOptions) error {
|
|
var paths []string
|
|
|
|
// Walk the directory tree and collect all paths
|
|
err := filepath.WalkDir(root, func(subPath string, d os.DirEntry, err error) error {
|
|
// Check for context cancellation
|
|
if ctx.Err() != nil {
|
|
return ctx.Err()
|
|
}
|
|
|
|
// Handle walk errors
|
|
if err != nil {
|
|
opts.Logger("[CLEAN] Walk error on %s: %v", subPath, err)
|
|
return nil // Skip errors and continue
|
|
}
|
|
|
|
// Skip the root directory itself
|
|
if subPath != root {
|
|
paths = append(paths, subPath)
|
|
}
|
|
return nil
|
|
})
|
|
|
|
// Process paths in reverse order (depth-first)
|
|
for i := len(paths) - 1; i >= 0; i-- {
|
|
select {
|
|
case <-ctx.Done():
|
|
return ctx.Err()
|
|
case workChan <- paths[i]:
|
|
// Path sent to worker
|
|
}
|
|
}
|
|
|
|
return err
|
|
}
|
|
|
|
// removeWorker processes paths from the work channel and removes them
|
|
func removeWorker(ctx context.Context, wg *sync.WaitGroup, workChan <-chan string, resultChan chan<- removeResult, opts RemoveOptions) {
|
|
defer wg.Done()
|
|
|
|
for {
|
|
select {
|
|
case <-ctx.Done():
|
|
return
|
|
case subPath, ok := <-workChan:
|
|
if !ok {
|
|
return // Channel closed, exit worker
|
|
}
|
|
|
|
// Try to remove with retries
|
|
var err error
|
|
for i := 0; i <= opts.Retries; i++ {
|
|
// Check if context is cancelled before each attempt
|
|
if ctx.Err() != nil {
|
|
resultChan <- removeResult{Path: subPath, Err: ctx.Err()}
|
|
return
|
|
}
|
|
|
|
err = os.RemoveAll(subPath)
|
|
if err == nil {
|
|
break // Successful removal
|
|
}
|
|
|
|
// If not the last retry, wait before trying again
|
|
if i < opts.Retries {
|
|
// Exponential backoff: 100ms, 200ms, 400ms, ...
|
|
backoff := time.Duration(100*(1<<i)) * time.Millisecond
|
|
select {
|
|
case <-ctx.Done():
|
|
resultChan <- removeResult{Path: subPath, Err: ctx.Err()}
|
|
return
|
|
case <-time.After(backoff):
|
|
// Continue with retry
|
|
}
|
|
}
|
|
}
|
|
|
|
// Send result
|
|
if err != nil {
|
|
opts.Logger("[CLEAN] Failed to remove %s: %v", subPath, err)
|
|
resultChan <- removeResult{Path: subPath, Err: fmt.Errorf("remove %q: %w", subPath, err)}
|
|
} else {
|
|
if opts.Logger != nil {
|
|
opts.Logger("[CLEAN] Removed: %s", subPath)
|
|
}
|
|
resultChan <- removeResult{Path: subPath, Err: nil}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// tryRemoveRoot attempts to remove the root directory itself
|
|
func tryRemoveRoot(ctx context.Context, path string, opts RemoveOptions) error {
|
|
// Check for context cancellation
|
|
select {
|
|
case <-ctx.Done():
|
|
return ctx.Err()
|
|
default:
|
|
// Continue
|
|
}
|
|
|
|
// Try to remove the directory
|
|
if err := os.Remove(path); err != nil {
|
|
// It's common for this to fail if the directory wasn't empty
|
|
opts.Logger("[CLEAN] Note: Could not remove root directory %s: %v", path, err)
|
|
return fmt.Errorf("remove root directory %q: %w", path, err)
|
|
}
|
|
|
|
opts.Logger("[CLEAN] Successfully removed root directory: %s", path)
|
|
return nil
|
|
}
|