github.com/ActiveState/cli@v0.0.0-20240508170324-6801f60cd051/scripts/ci/parallelize/parallelize.go (about)

     1  package main
     2  
     3  import (
     4  	"encoding/json"
     5  	"fmt"
     6  	"os"
     7  	"os/exec"
     8  	"path/filepath"
     9  	"strconv"
    10  	"strings"
    11  	"time"
    12  
    13  	"github.com/ActiveState/cli/internal/constraints"
    14  	"github.com/ActiveState/cli/internal/errs"
    15  	"github.com/ActiveState/cli/internal/fileutils"
    16  	"github.com/ActiveState/cli/internal/installation/storage"
    17  	"github.com/ActiveState/cli/internal/osutils"
    18  	"github.com/ActiveState/cli/pkg/project"
    19  	"github.com/gammazero/workerpool"
    20  )
    21  
    22  type Job struct {
    23  	ID   string
    24  	Args []string
    25  	Env  []string
    26  	If   string
    27  }
    28  
    29  func main() {
    30  	if err := run(); err != nil {
    31  		fmt.Fprintln(os.Stderr, errs.JoinMessage(err))
    32  		os.Exit(1)
    33  	}
    34  }
    35  
    36  func run() error {
    37  	if len(os.Args) <= 1 {
    38  		return errs.New("Must provide single argument with JSON blob, or [job <ID>] to check the results of a job.")
    39  	}
    40  
    41  	if os.Args[1] == "results" {
    42  		if len(os.Args) != 3 {
    43  			return errs.New("Must provide job ID")
    44  		}
    45  		return readJob(os.Args[2])
    46  	}
    47  
    48  	jsonData := []byte(strings.Join(os.Args[1:], ""))
    49  	var jobs []Job
    50  	if err := json.Unmarshal(jsonData, &jobs); err != nil {
    51  		return errs.Wrap(err, "Invalid JSON. Data: %s", jsonData)
    52  	}
    53  
    54  	wp := workerpool.New(3)
    55  	for _, job := range jobs {
    56  		func(job Job) {
    57  			wp.Submit(func() {
    58  				t := time.Now()
    59  				fmt.Printf("Running: %s\n", job.ID)
    60  				runJob(job)
    61  				fmt.Printf("Finished %s after %s\n", job.ID, time.Since(t))
    62  			})
    63  		}(job)
    64  	}
    65  	wp.StopWait()
    66  
    67  	return nil
    68  }
    69  
    70  func jobDir() string {
    71  	path, err := storage.AppDataPath()
    72  	if err != nil {
    73  		panic(err)
    74  	}
    75  
    76  	path = filepath.Join(path, "jobs")
    77  	if err := fileutils.MkdirUnlessExists(path); err != nil {
    78  		panic(err)
    79  	}
    80  
    81  	return path
    82  }
    83  
    84  func runJob(job Job) {
    85  	outname := filepath.Join(jobDir(), fmt.Sprintf("%s.out", job.ID))
    86  	fmt.Printf("%s: saving to %s\n", job.ID, outname)
    87  
    88  	outfile, err := os.Create(outname)
    89  	if err != nil {
    90  		panic(fmt.Sprintf("Could not create: %#v, error: %s\n", job, errs.JoinMessage(err)))
    91  	}
    92  	defer outfile.Close()
    93  
    94  	failure := func(msg string, args ...interface{}) {
    95  		fmt.Fprintf(outfile, msg+"\n1", args...)
    96  		fmt.Fprintf(os.Stderr, fmt.Sprintf("%s: ", job.ID)+msg, args...)
    97  	}
    98  
    99  	if job.If != "" {
   100  		pj, err := project.GetOnce()
   101  		if err != nil {
   102  			failure("Could not get project: %s", errs.JoinMessage(err))
   103  			return
   104  		}
   105  
   106  		cond := constraints.NewPrimeConditional(nil, pj, "")
   107  		run, err := cond.Eval(job.If)
   108  		if err != nil {
   109  			failure("Could not evaluate conditonal: %s, error: %s\n", job.If, errs.JoinMessage(err))
   110  			return
   111  		}
   112  		if !run {
   113  			fmt.Printf("%s: Skipping as per conditional: %s\n", job.ID, job.If)
   114  			return
   115  		}
   116  	}
   117  	if len(job.Args) == 0 {
   118  		failure("Job must have arguments: %#v\n", job)
   119  		return
   120  	}
   121  
   122  	code, _, err := osutils.Execute(job.Args[0]+osutils.ExeExtension, job.Args[1:], func(cmd *exec.Cmd) error {
   123  		cmd.Stdout = outfile
   124  		cmd.Stderr = outfile
   125  		cmd.Env = append(job.Env, os.Environ()...)
   126  		return nil
   127  	})
   128  	if err != nil {
   129  		failure("Executing job %s failed, error: %s", job.ID, errs.JoinMessage(err))
   130  	}
   131  	_, err = outfile.WriteString(fmt.Sprintf("\n%d", code)) // last entry is the exit code
   132  	if err != nil {
   133  		fmt.Printf("Could not write exit code to file: %s\n", errs.JoinMessage(err))
   134  	}
   135  	fmt.Printf("%s: Completed, exit code: %d\n", job.ID, code)
   136  }
   137  
   138  func readJob(id string) error {
   139  	jobfile := filepath.Join(jobDir(), fmt.Sprintf("%s.out", id))
   140  	if !fileutils.FileExists(jobfile) {
   141  		return errs.New("Job does not exist: %s", jobfile)
   142  	}
   143  
   144  	contents := strings.Split(string(fileutils.ReadFileUnsafe(jobfile)), "\n")
   145  	code, err := strconv.Atoi(contents[len(contents)-1])
   146  	if err != nil {
   147  		return errs.Wrap(err, "Expected last line to be the exit code, instead found: %s", contents[len(contents)-1])
   148  	}
   149  
   150  	fmt.Println(strings.Join(contents[0:(len(contents)-2)], "\n"))
   151  	os.Exit(code)
   152  
   153  	return nil
   154  }