github.com/nilium/gitlab-runner@v12.5.0+incompatible/scripts/pull-images-for-tests/main.go (about)

     1  package main
     2  
     3  import (
     4  	"bufio"
     5  	"context"
     6  	"fmt"
     7  	"io"
     8  	"os"
     9  	"os/exec"
    10  	"sync"
    11  
    12  	"gitlab.com/gitlab-org/gitlab-runner/common"
    13  	"gitlab.com/gitlab-org/gitlab-runner/helpers"
    14  )
    15  
    16  var images = []string{
    17  	common.TestAlpineImage,
    18  	common.TestAlpineNoRootImage,
    19  	common.TestDockerGitImage,
    20  	common.TestDockerDindImage,
    21  }
    22  
    23  type rewriter struct {
    24  	ctx    context.Context
    25  	prefix string
    26  	input  *bufio.Reader
    27  }
    28  
    29  func (rw *rewriter) watch() {
    30  	for {
    31  		select {
    32  		case <-rw.ctx.Done():
    33  			return
    34  		case err := <-rw.rewriteInput():
    35  			if err != nil {
    36  				rw.writeToOutput(fmt.Sprintf("Error while reading command output: %v", err))
    37  				return
    38  			}
    39  		}
    40  	}
    41  }
    42  
    43  func (rw *rewriter) rewriteInput() <-chan error {
    44  	e := make(chan error)
    45  
    46  	go func() {
    47  		line, err := rw.input.ReadString('\n')
    48  		if err == nil || err == io.EOF {
    49  			rw.writeToOutput(line)
    50  			e <- nil
    51  
    52  			return
    53  		}
    54  
    55  		e <- err
    56  	}()
    57  
    58  	return e
    59  }
    60  
    61  func (rw *rewriter) writeToOutput(line string) {
    62  	fmt.Printf("%s[%s]%s %s", helpers.ANSI_YELLOW, rw.prefix, helpers.ANSI_RESET, line)
    63  }
    64  
    65  func newRewriter(ctx context.Context, prefix string) io.Writer {
    66  	pr, pw, err := os.Pipe()
    67  	if err != nil {
    68  		panic(err)
    69  	}
    70  
    71  	w := &rewriter{
    72  		ctx:    ctx,
    73  		prefix: prefix,
    74  		input:  bufio.NewReader(pr),
    75  	}
    76  
    77  	go w.watch()
    78  
    79  	return pw
    80  }
    81  
    82  func pullImage(wg *sync.WaitGroup, name string) {
    83  	ctx, cancel := context.WithCancel(context.Background())
    84  	defer func() {
    85  		cancel()
    86  		wg.Done()
    87  	}()
    88  
    89  	output := newRewriter(ctx, name)
    90  
    91  	cmd := exec.Command("docker", "pull", name)
    92  	cmd.Stdout = output
    93  	cmd.Stderr = output
    94  
    95  	err := cmd.Run()
    96  	if err != nil {
    97  		panic(err)
    98  	}
    99  }
   100  
   101  func main() {
   102  	wg := new(sync.WaitGroup)
   103  
   104  	for _, image := range images {
   105  		wg.Add(1)
   106  		go pullImage(wg, image)
   107  	}
   108  
   109  	wg.Wait()
   110  }