feat: context propagation

Signed-off-by: Ruslan Semagin <[email protected]>

cmd/build.go

@@ -81,6 +81,8 @@ func build(cmd *cobra.Command, _ []string) error {
 		return err
 	}
 
+	module.Ctx = cmd.Context()
+
 	if err := module.Build(); err != nil {
 		return err
 	}

internal/fs.go

@@ -1,6 +1,7 @@
 package internal
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"io"
@@ -13,6 +14,7 @@ import (
 )
 
 func copyFromPath(
+	ctx context.Context,
 	wg *sync.WaitGroup,
 	errCh chan<- error,
 	ignore *[]string,
@@ -21,14 +23,20 @@ func copyFromPath(
 ) {
 	defer wg.Done()
 
-	if err := walk(wg, errCh, from, to, ignore, existsMode); err != nil {
+	if err := CheckContextActivity(ctx); err != nil {
+		errCh <- err
+		return
+	}
+
+	if err := walk(ctx, wg, errCh, from, to, ignore, existsMode); err != nil {
 		if !errors.Is(err, doublestar.SkipDir) {
 			errCh <- err
 		}
 	}
 }
 
 func walk(
+	ctx context.Context,
 	wg *sync.WaitGroup,
 	errCh chan<- error,
 	from, to string,
@@ -42,6 +50,11 @@ func walk(
 	jobs := make(chan struct{}, 10)
 
 	err := filepath.Walk(from, func(path string, info os.FileInfo, err error) error {
+		if ctxErr := CheckContextActivity(ctx); ctxErr != nil {
+			errCh <- ctxErr
+			return ctxErr
+		}
+
 		if err != nil {
 			return err
 		}
@@ -82,7 +95,7 @@ func walk(
 			}
 
 			wg2.Add(1)
-			go copyFile(&wg2, errCh, absFrom, absTo, jobs, existsMode)
+			go copyFile(ctx, &wg2, errCh, absFrom, absTo, jobs, existsMode)
 		}
 
 		return nil
@@ -94,6 +107,7 @@ func walk(
 }
 
 func copyFile(
+	ctx context.Context,
 	wg *sync.WaitGroup,
 	errCh chan<- error,
 	src, dst string,
@@ -102,6 +116,11 @@ func copyFile(
 ) {
 	defer wg.Done()
 
+	if err := CheckContextActivity(ctx); err != nil {
+		errCh <- err
+		return
+	}
+
 	fileName := strings.LastIndex(src, "/")
 	if !strings.HasSuffix(dst, src[fileName:]) {
 		dst = filepath.Join(dst, src[fileName:])
@@ -133,6 +152,11 @@ func copyFile(
 		}
 	}(in)
 
+	if err := CheckContextActivity(ctx); err != nil {
+		errCh <- err
+		return
+	}
+
 	info, err := in.Stat()
 	if err != nil {
 		errCh <- err
@@ -160,12 +184,22 @@ func copyFile(
 			}
 		}(out)
 
+		if err := CheckContextActivity(ctx); err != nil {
+			errCh <- err
+			return
+		}
+
 		_, err = io.Copy(out, in)
 		if err != nil {
 			errCh <- err
 			return
 		}
 
+		if err := CheckContextActivity(ctx); err != nil {
+			errCh <- err
+			return
+		}
+
 		err = os.Chtimes(dst, info.ModTime(), info.ModTime())
 		if err != nil {
 			errCh <- err

internal/module.go

@@ -1,6 +1,7 @@
 package internal
 
 import (
+	"context"
 	"errors"
 	"fmt"
 )
@@ -21,6 +22,7 @@ type Item struct {
 }
 
 type Module struct {
+	Ctx            context.Context
 	Name           string   `yaml:"name"`
 	Version        string   `yaml:"version"`
 	Account        string   `yaml:"account"`

internal/module_builder.go

@@ -1,6 +1,7 @@
 package internal
 
 import (
+	"context"
 	"fmt"
 	"os"
 	"path/filepath"
@@ -11,6 +12,10 @@ import (
 )
 
 func (m *Module) Build() error {
+	if err := CheckContextActivity(m.Ctx); err != nil {
+		return err
+	}
+
 	logFile, err := os.OpenFile(
 		fmt.Sprintf("./%s-%s.%s.log", m.Name, m.Version, time.Now().UTC().Format(time.RFC3339)),
 		os.O_CREATE|os.O_WRONLY|os.O_APPEND,
@@ -131,8 +136,12 @@ func (m *Module) Collect(log *zerolog.Logger) error {
 	errCh := make(chan error, len(m.Mapping))
 
 	for _, item := range m.Mapping {
+		if err := CheckContextActivity(m.Ctx); err != nil {
+			return err
+		}
+
 		wg.Add(1)
-		go handleItem(&wg, errCh, &m.Ignore, item, buildDirectory)
+		go handleItem(m.Ctx, &wg, errCh, &m.Ignore, item, buildDirectory)
 	}
 
 	wg.Wait()
@@ -152,6 +161,7 @@ func (m *Module) Collect(log *zerolog.Logger) error {
 }
 
 func handleItem(
+	ctx context.Context,
 	wg *sync.WaitGroup,
 	errCh chan<- error,
 	ignore *[]string,
@@ -160,14 +170,24 @@ func handleItem(
 ) {
 	defer wg.Done()
 
+	if err := CheckContextActivity(ctx); err != nil {
+		errCh <- err
+		return
+	}
+
 	to, err := mkdir(fmt.Sprintf("%s/%s", buildDirectory, item.RelativePath))
 	if err != nil {
 		errCh <- err
 		return
 	}
 
 	for _, from := range item.Paths {
+		if err := CheckContextActivity(ctx); err != nil {
+			errCh <- err
+			return
+		}
+
 		wg.Add(1)
-		go copyFromPath(wg, errCh, ignore, from, to, item.IfFileExists)
+		go copyFromPath(ctx, wg, errCh, ignore, from, to, item.IfFileExists)
 	}
 }

internal/validators.go

@@ -1,6 +1,7 @@
 package internal
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"os"
@@ -38,3 +39,12 @@ func ValidateVersion(version string) error {
 
 	return fmt.Errorf("invalid module version %s", version)
 }
+
+func CheckContextActivity(ctx context.Context) error {
+	select {
+	case <-ctx.Done():
+		return fmt.Errorf("context canceled: %w", ctx.Err())
+	default:
+		return nil
+	}
+}