Cleanup logic for getting a directory size; support cancelation during goroutine

server/environment_docker.go

@@ -146,6 +146,10 @@ func (d *DockerEnvironment) OnBeforeStart() error {
 		return err
 	}
 
+	if !d.Server.Filesystem.HasSpaceAvailable() {
+		return errors.New("cannot start server, not enough disk space available")
+	}
+
 	// Always destroy and re-create the server container to ensure that synced data from
 	// the Panel is used.
 	if err := d.Client.ContainerRemove(context.Background(), d.Server.Uuid, types.ContainerRemoveOptions{RemoveVolumes: true}); err != nil {

server/filesystem.go

@@ -3,6 +3,7 @@ package server
 import (
 	"bufio"
 	"bytes"
+	"context"
 	"encoding/json"
 	"fmt"
 	"github.com/gabriel-vasile/mimetype"
@@ -67,7 +68,7 @@ func (fs *Filesystem) SafePath(p string) (string, error) {
 		// Range over all of the path parts and form directory pathings from the end
 		// moving up until we have a valid resolution or we run out of paths to try.
 		for k := range parts {
-			try = strings.Join(parts[:(len(parts) - k)], "/")
+			try = strings.Join(parts[:(len(parts)-k)], "/")
 
 			if !strings.HasPrefix(try, fs.Path()) {
 				break
@@ -118,23 +119,27 @@ func (fs *Filesystem) HasSpaceAvailable() bool {
 		return true
 	}
 
-	var size int64
+	// If we have a match in the cache, use that value in the return. No need to perform an expensive
+	// disk operation, even if this is an empty value.
 	if x, exists := fs.Server.Cache.Get("disk_used"); exists {
-		size = x.(int64)
+		fs.Server.Resources.Disk = x.(int64)
+		return (x.(int64) / 1000.0 / 1000.0) <= space
 	}
 
 	// If there is no size its either because there is no data (in which case running this function
 	// will have effectively no impact), or there is nothing in the cache, in which case we need to
 	// grab the size of their data directory. This is a taxing operation, so we want to store it in
 	// the cache once we've gotten it.
-	if size == 0 {
-		if size, err := fs.DirectorySize("/"); err != nil {
-			zap.S().Warnw("failed to determine directory size", zap.String("server", fs.Server.Uuid), zap.Error(err))
-		} else {
-			fs.Server.Cache.Set("disk_used", size, time.Second * 60)
-		}
+	size, err := fs.DirectorySize("/")
+	if err != nil {
+		zap.S().Warnw("failed to determine directory size", zap.String("server", fs.Server.Uuid), zap.Error(err))
 	}
 
+	// Always cache the size, even if there is an error. We want to always return that value
+	// so that we don't cause an endless loop of determining the disk size if there is a temporary
+	// error encountered.
+	fs.Server.Cache.Set("disk_used", size, time.Second*60)
+
 	// Determine if their folder size, in bytes, is smaller than the amount of space they've
 	// been allocated.
 	fs.Server.Resources.Disk = size
@@ -146,42 +151,15 @@ func (fs *Filesystem) HasSpaceAvailable() bool {
 // through all of the folders. Returns the size in bytes. This can be a fairly taxing operation
 // on locations with tons of files, so it is recommended that you cache the output.
 func (fs *Filesystem) DirectorySize(dir string) (int64, error) {
+	w := fs.NewWalker()
+	ctx := context.Background()
+
 	var size int64
-	var wg sync.WaitGroup
+	err := w.Walk(dir, ctx, func(f os.FileInfo) {
+		atomic.AddInt64(&size, f.Size())
+	})
 
-	cleaned, err := fs.SafePath(dir)
-	if err != nil {
-		return 0, err
-	}
-
-	files, err := ioutil.ReadDir(cleaned)
-	if err != nil {
-		return 0, err
-	}
-
-	// Iterate over all of the files and directories. If it is a file, immediately add its size
-	// to the total size being returned. If we're dealing with a directory, call this function
-	// on a seperate thread until we have gotten the size of everything nested within the given
-	// directory.
-	for _, f := range files {
-		if f.IsDir() {
-			wg.Add(1)
-
-			go func(p string) {
-				defer wg.Done()
-
-				s, _ := fs.DirectorySize(p)
-
-				atomic.AddInt64(&size, s)
-			}(filepath.Join(cleaned, f.Name()))
-		} else {
-			atomic.AddInt64(&size, f.Size())
-		}
-	}
-
-	wg.Wait()
-
-	return size, nil
+	return size, err
 }
 
 // Reads a file on the system and returns it as a byte representation in a file

server/filesystem_walker.go

@@ -0,0 +1,61 @@
+package server
+
+import (
+	"context"
+	"golang.org/x/sync/errgroup"
+	"io/ioutil"
+	"os"
+	"path/filepath"
+)
+
+type FileWalker struct {
+	*Filesystem
+}
+
+// Returns a new walker instance.
+func (fs *Filesystem) NewWalker() *FileWalker {
+	return &FileWalker{fs}
+}
+
+// Iterate over all of the files and directories within a given directory. When a file is
+// found the callback will be called with the file information. If a directory is encountered
+// it will be recursively passed back through to this function.
+func (fw *FileWalker) Walk(dir string, ctx context.Context, callback func (os.FileInfo)) error {
+	cleaned, err := fw.SafePath(dir)
+	if err != nil {
+		return err
+	}
+
+	// Get all of the files from this directory.
+	files, err := ioutil.ReadDir(cleaned)
+	if err != nil {
+		return err
+	}
+
+	// Create an error group that we can use to run processes in parallel while retaining
+	// the ability to cancel the entire process immediately should any of it fail.
+	g, ctx := errgroup.WithContext(ctx)
+
+	for _, f := range files {
+		if f.IsDir() {
+			p := filepath.Join(dir, f.Name())
+			// Recursively call this function to continue digging through the directory tree within
+			// a seperate goroutine. If the context is canceled abort this process.
+			g.Go(func() error {
+				select {
+				case <-ctx.Done():
+					return ctx.Err()
+				default:
+					return fw.Walk(p, ctx, callback)
+				}
+			})
+		} else {
+			// If this isn't a directory, go ahead and pass the file information into the
+			// callback.
+			callback(f)
+		}
+	}
+
+	// Block until all of the routines finish and have returned a value.
+	return g.Wait()
+}