github.com/abayer/test-infra@v0.0.5/prow/cmd/tot/main.go (about) 1 /* 2 Copyright 2016 The Kubernetes Authors. 3 4 Licensed under the Apache License, Version 2.0 (the "License"); 5 you may not use this file except in compliance with the License. 6 You may obtain a copy of the License at 7 8 http://www.apache.org/licenses/LICENSE-2.0 9 10 Unless required by applicable law or agreed to in writing, software 11 distributed under the License is distributed on an "AS IS" BASIS, 12 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 See the License for the specific language governing permissions and 14 limitations under the License. 15 */ 16 17 // Tot vends (rations) incrementing numbers for use in builds. 18 // https://en.wikipedia.org/wiki/Rum_ration 19 package main 20 21 import ( 22 "encoding/json" 23 "errors" 24 "flag" 25 "fmt" 26 "io/ioutil" 27 "net/http" 28 "os" 29 "strconv" 30 "strings" 31 "sync" 32 "time" 33 34 "github.com/sirupsen/logrus" 35 36 "k8s.io/test-infra/prow/config" 37 "k8s.io/test-infra/prow/logrusutil" 38 "k8s.io/test-infra/prow/pjutil" 39 "k8s.io/test-infra/prow/pod-utils/downwardapi" 40 "k8s.io/test-infra/prow/pod-utils/gcs" 41 ) 42 43 type options struct { 44 port int 45 storagePath string 46 47 useFallback bool 48 fallbackURI string 49 50 configPath string 51 jobConfigPath string 52 fallbackBucket string 53 } 54 55 func gatherOptions() options { 56 o := options{} 57 flag.IntVar(&o.port, "port", 8888, "Port to listen on.") 58 flag.StringVar(&o.storagePath, "storage", "tot.json", "Where to store the results.") 59 60 flag.BoolVar(&o.useFallback, "fallback", false, "Fallback to GCS bucket for missing builds.") 61 flag.StringVar(&o.fallbackURI, "fallback-url-template", 62 "https://storage.googleapis.com/kubernetes-jenkins/logs/%s/latest-build.txt", 63 "URL template to fallback to for jobs that lack a last vended build number.", 64 ) 65 66 flag.StringVar(&o.configPath, "config-path", "", "Path to prow config.") 67 flag.StringVar(&o.jobConfigPath, "job-config-path", "", "Path to prow job configs.") 68 flag.StringVar(&o.fallbackBucket, "fallback-bucket", "", 69 "Fallback to top-level bucket for jobs that lack a last vended build number. The bucket layout is expected to follow https://github.com/kubernetes/test-infra/tree/master/gubernator#gcs-bucket-layout", 70 ) 71 72 flag.Parse() 73 return o 74 } 75 76 func (o *options) Validate() error { 77 if o.configPath != "" && o.fallbackBucket == "" { 78 return errors.New("you need to provide a bucket to fallback to when the prow config is specified") 79 } 80 if o.configPath == "" && o.fallbackBucket != "" { 81 return errors.New("you need to provide the prow config when a fallback bucket is specified") 82 } 83 return nil 84 } 85 86 type store struct { 87 Number map[string]int // job name -> last vended build number 88 mutex sync.Mutex 89 storagePath string 90 fallbackFunc func(string) int 91 } 92 93 func newStore(storagePath string) (*store, error) { 94 s := &store{ 95 Number: make(map[string]int), 96 storagePath: storagePath, 97 } 98 buf, err := ioutil.ReadFile(storagePath) 99 if err == nil { 100 err = json.Unmarshal(buf, s) 101 if err != nil { 102 return nil, err 103 } 104 } else if !os.IsNotExist(err) { 105 return nil, err 106 } 107 return s, nil 108 } 109 110 func (s *store) save() error { 111 buf, err := json.Marshal(s) 112 if err != nil { 113 return err 114 } 115 err = ioutil.WriteFile(s.storagePath+".tmp", buf, 0644) 116 if err != nil { 117 return err 118 } 119 return os.Rename(s.storagePath+".tmp", s.storagePath) 120 } 121 122 func (s *store) vend(jobName string) int { 123 s.mutex.Lock() 124 defer s.mutex.Unlock() 125 n, ok := s.Number[jobName] 126 if !ok && s.fallbackFunc != nil { 127 n = s.fallbackFunc(jobName) 128 } 129 n++ 130 131 s.Number[jobName] = n 132 133 err := s.save() 134 if err != nil { 135 logrus.Error(err) 136 } 137 138 return n 139 } 140 141 func (s *store) peek(jobName string) int { 142 s.mutex.Lock() 143 defer s.mutex.Unlock() 144 return s.Number[jobName] 145 } 146 147 func (s *store) set(jobName string, n int) { 148 s.mutex.Lock() 149 defer s.mutex.Unlock() 150 s.Number[jobName] = n 151 152 err := s.save() 153 if err != nil { 154 logrus.Error(err) 155 } 156 } 157 158 func (s *store) handle(w http.ResponseWriter, r *http.Request) { 159 jobName := r.URL.Path[len("/vend/"):] 160 switch r.Method { 161 case "GET": 162 n := s.vend(jobName) 163 logrus.Infof("Vending %s number %d to %s.", jobName, n, r.RemoteAddr) 164 fmt.Fprintf(w, "%d", n) 165 case "HEAD": 166 n := s.peek(jobName) 167 logrus.Infof("Peeking %s number %d to %s.", jobName, n, r.RemoteAddr) 168 fmt.Fprintf(w, "%d", n) 169 case "POST": 170 body, err := ioutil.ReadAll(r.Body) 171 if err != nil { 172 logrus.WithError(err).Error("Unable to read body.") 173 return 174 } 175 n, err := strconv.Atoi(string(body)) 176 if err != nil { 177 logrus.WithError(err).Error("Unable to parse number.") 178 return 179 } 180 logrus.Infof("Setting %s to %d from %s.", jobName, n, r.RemoteAddr) 181 s.set(jobName, n) 182 } 183 } 184 185 type fallbackHandler struct { 186 template string 187 // in case a config agent is provided, tot will 188 // determine the GCS path that it needs to use 189 // based on the configured jobs in prow and 190 // bucket. 191 configAgent *config.Agent 192 bucket string 193 } 194 195 func (f fallbackHandler) get(jobName string) int { 196 url := f.getURL(jobName) 197 198 var body []byte 199 200 for i := 0; i < 10; i++ { 201 resp, err := http.Get(url) 202 if err == nil { 203 defer resp.Body.Close() 204 if resp.StatusCode == http.StatusOK { 205 body, err = ioutil.ReadAll(resp.Body) 206 if err == nil { 207 break 208 } else { 209 logrus.WithError(err).Error("Failed to read response body.") 210 } 211 } else if resp.StatusCode == http.StatusNotFound { 212 break 213 } 214 } else { 215 logrus.WithError(err).Errorf("Failed to GET %s.", url) 216 } 217 time.Sleep(2 * time.Second) 218 } 219 220 n, err := strconv.Atoi(strings.TrimSpace(string(body))) 221 if err != nil { 222 return 0 223 } 224 225 return n 226 } 227 228 func (f fallbackHandler) getURL(jobName string) string { 229 if f.configAgent == nil { 230 return fmt.Sprintf(f.template, jobName) 231 } 232 233 var spec *downwardapi.JobSpec 234 cfg := f.configAgent.Config() 235 236 for _, pre := range cfg.AllPresubmits(nil) { 237 if jobName == pre.Name { 238 spec = pjutil.PresubmitToJobSpec(pre) 239 break 240 } 241 } 242 if spec == nil { 243 for _, post := range cfg.AllPostsubmits(nil) { 244 if jobName == post.Name { 245 spec = pjutil.PostsubmitToJobSpec(post) 246 break 247 } 248 } 249 } 250 if spec == nil { 251 for _, per := range cfg.AllPeriodics() { 252 if jobName == per.Name { 253 spec = pjutil.PeriodicToJobSpec(per) 254 break 255 } 256 } 257 } 258 // If spec is still nil, we know nothing about the requested job. 259 if spec == nil { 260 logrus.Errorf("requested job is unknown to prow: %s", jobName) 261 return "" 262 } 263 paths := gcs.LatestBuildForSpec(spec, nil) 264 if len(paths) != 1 { 265 logrus.Errorf("expected a single GCS path, got %v", paths) 266 return "" 267 } 268 return fmt.Sprintf("%s/%s", strings.TrimSuffix(f.bucket, "/"), paths[0]) 269 } 270 271 func main() { 272 o := gatherOptions() 273 if err := o.Validate(); err != nil { 274 logrus.Fatalf("Invalid options: %v", err) 275 } 276 logrus.SetFormatter( 277 logrusutil.NewDefaultFieldsFormatter(nil, logrus.Fields{"component": "tot"}), 278 ) 279 280 s, err := newStore(o.storagePath) 281 if err != nil { 282 logrus.WithError(err).Fatal("newStore failed") 283 } 284 285 if o.useFallback { 286 var configAgent *config.Agent 287 if o.configPath != "" { 288 configAgent = &config.Agent{} 289 if err := configAgent.Start(o.configPath, o.jobConfigPath); err != nil { 290 logrus.WithError(err).Fatal("Error starting config agent.") 291 } 292 } 293 294 s.fallbackFunc = fallbackHandler{ 295 template: o.fallbackURI, 296 configAgent: configAgent, 297 bucket: o.fallbackBucket, 298 }.get 299 } 300 301 http.HandleFunc("/vend/", s.handle) 302 303 logrus.Fatal(http.ListenAndServe(":"+strconv.Itoa(o.port), nil)) 304 }