kraken

1
// Copyright (c) 2016-2019 Uber Technologies, Inc.
2
//
3
// Licensed under the Apache License, Version 2.0 (the "License");
4
// you may not use this file except in compliance with the License.
5
// You may obtain a copy of the License at
6
//
7
//     http://www.apache.org/licenses/LICENSE-2.0
8
//
9
// Unless required by applicable law or agreed to in writing, software
10
// distributed under the License is distributed on an "AS IS" BASIS,
11
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
// See the License for the specific language governing permissions and
13
// limitations under the License.
14
package main
15

16
import (
17
	"bytes"
18
	"encoding/json"
19
	"flag"
20
	"fmt"
21
	"os"
22
	"strings"
23
	"sync"
24
	"time"
25

26
	"github.com/uber/kraken/lib/torrent/scheduler"
27
	"github.com/uber/kraken/utils/configutil"
28
	"github.com/uber/kraken/utils/httputil"
29
	"github.com/uber/kraken/utils/osutil"
30
)
31

32
type appConfig struct {
33
	Scheduler scheduler.Config `yaml:"scheduler"`
34
}
35

36
func reload(addr string, config scheduler.Config) error {
37
	b, err := json.Marshal(config)
38
	if err != nil {
39
		return err
40
	}
41
	_, err = httputil.Patch(
42
		fmt.Sprintf("http://%s/x/config/scheduler", addr),
43
		httputil.SendBody(bytes.NewBuffer(b)),
44
		httputil.SendTimeout(5*time.Second))
45
	return err
46
}
47

48
func main() {
49
	configFile := flag.String("config", "", "config file")
50
	hostFile := flag.String("f", "", "host file")
51
	hostStr := flag.String("hosts", "", "comma-separated hosts")
52
	port := flag.Int("port", 7602, "server port (different for agent / origin)")
53
	flag.Parse()
54

55
	if *configFile == "" {
56
		panic("-config required")
57
	}
58
	if (*hostFile != "" && *hostStr != "") || (*hostFile == "" && *hostStr == "") {
59
		panic("must set either -f or -hosts")
60
	}
61
	if *port == 0 {
62
		panic("-port must be non-zero")
63
	}
64

65
	var hosts []string
66
	if *hostFile != "" {
67
		f, err := os.Open(*hostFile)
68
		if err != nil {
69
			panic(err)
70
		}
71
		hosts, err = osutil.ReadLines(f)
72
		if err != nil {
73
			panic(err)
74
		}
75
	} else if *hostStr != "" {
76
		hosts = strings.Split(*hostStr, ",")
77
	}
78

79
	var config appConfig
80
	if err := configutil.Load(*configFile, &config); err != nil {
81
		panic(err)
82
	}
83

84
	errs := make(chan error)
85
	var wg sync.WaitGroup
86
	for _, host := range hosts {
87
		wg.Add(1)
88
		go func(host string) {
89
			defer wg.Done()
90
			addr := fmt.Sprintf("%s:%d", host, *port)
91
			if err := reload(addr, config.Scheduler); err != nil {
92
				errs <- err
93
			}
94
		}(host)
95
	}
96
	go func() {
97
		for err := range errs {
98
			fmt.Println(err)
99
		}
100
	}()
101
	wg.Wait()
102
	close(errs)
103
}
104