1package main
2
3import (
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
16var images = []string{
17	common.TestAlpineImage,
18	common.TestAlpineNoRootImage,
19	common.TestDockerGitImage,
20	common.TestDockerDindImage,
21}
22
23type rewriter struct {
24	ctx    context.Context
25	prefix string
26	input  *bufio.Reader
27}
28
29func (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
43func (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
61func (rw *rewriter) writeToOutput(line string) {
62	fmt.Printf("%s[%s]%s %s", helpers.ANSI_YELLOW, rw.prefix, helpers.ANSI_RESET, line)
63}
64
65func 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
82func 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
101func 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}
111