github.com/munnerz/test-infra@v0.0.0-20190108210205-ce3d181dc989/boskos/janitor/janitor.go (about)

     1  /*
     2  Copyright 2017 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  package main
    18  
    19  import (
    20  	"flag"
    21  	"fmt"
    22  	"os/exec"
    23  	"time"
    24  
    25  	"github.com/sirupsen/logrus"
    26  	"k8s.io/test-infra/boskos/client"
    27  	"k8s.io/test-infra/boskos/common"
    28  )
    29  
    30  var (
    31  	bufferSize     = 1 // Maximum holding resources
    32  	serviceAccount = flag.String("service-account", "", "Path to projects service account")
    33  	rTypes         common.CommaSeparatedStrings
    34  	poolSize       int
    35  	janitorPath    = flag.String("janitor-path", "/bin/janitor.py", "Path to janitor binary path")
    36  	boskosURL      = flag.String("boskos-url", "http://boskos", "Boskos URL")
    37  )
    38  
    39  func init() {
    40  	flag.Var(&rTypes, "resource-type", "comma-separated list of resources need to be cleaned up")
    41  	flag.IntVar(&poolSize, "pool-size", 20, "number of concurrent janitor goroutine")
    42  }
    43  
    44  func main() {
    45  	// Activate service account
    46  	flag.Parse()
    47  
    48  	logrus.SetFormatter(&logrus.JSONFormatter{})
    49  	boskos := client.NewClient("Janitor", *boskosURL)
    50  	logrus.Info("Initialized boskos client!")
    51  
    52  	if *serviceAccount == "" {
    53  		logrus.Fatal("--service-account cannot be empty!")
    54  	}
    55  
    56  	if len(rTypes) == 0 {
    57  		logrus.Fatal("--resource-type must not be empty!")
    58  	}
    59  
    60  	cmd := exec.Command("gcloud", "auth", "activate-service-account", "--key-file="+*serviceAccount)
    61  	if b, err := cmd.CombinedOutput(); err != nil {
    62  		logrus.WithError(err).Fatalf("fail to activate service account from %s :%s", *serviceAccount, string(b))
    63  	}
    64  
    65  	buffer := setup(boskos, poolSize, bufferSize, janitorClean)
    66  
    67  	for {
    68  		run(boskos, buffer, rTypes)
    69  		time.Sleep(time.Minute)
    70  	}
    71  }
    72  
    73  type clean func(string) error
    74  
    75  // Clean by janitor script
    76  func janitorClean(proj string) error {
    77  	cmd := exec.Command(*janitorPath, fmt.Sprintf("--project=%s", proj), "--hour=0")
    78  	b, err := cmd.CombinedOutput()
    79  	if err != nil {
    80  		logrus.WithError(err).Errorf("failed to clean up project %s, error info: %s", proj, string(b))
    81  	} else {
    82  		logrus.Infof("successfully cleaned up project %s", proj)
    83  	}
    84  	return err
    85  }
    86  
    87  type boskosClient interface {
    88  	Acquire(rtype string, state string, dest string) (*common.Resource, error)
    89  	ReleaseOne(name string, dest string) error
    90  }
    91  
    92  func setup(c boskosClient, janitorCount int, bufferSize int, cleanFunc clean) chan string {
    93  	buffer := make(chan string, 1)
    94  	for i := 0; i < janitorCount; i++ {
    95  		go janitor(c, buffer, cleanFunc)
    96  	}
    97  	return buffer
    98  }
    99  
   100  func run(c boskosClient, buffer chan string, rtypes []string) int {
   101  	totalAcquire := 0
   102  	res := make(map[string]int)
   103  	for _, s := range rtypes {
   104  		res[s] = 0
   105  	}
   106  
   107  	for {
   108  		for r := range res {
   109  			if projRes, err := c.Acquire(r, common.Dirty, common.Cleaning); err != nil {
   110  				logrus.WithError(err).Error("boskos acquire failed!")
   111  				totalAcquire += res[r]
   112  				delete(res, r)
   113  			} else if projRes == nil {
   114  				// To Sen: I don t understand why this would happen
   115  				logrus.Warning("received nil resource")
   116  				totalAcquire += res[r]
   117  				delete(res, r)
   118  			} else {
   119  				logrus.Infof("Acquired resources %s of type %s", projRes.Name, projRes.Type)
   120  				buffer <- projRes.Name // will block until buffer has a free slot
   121  				res[r]++
   122  			}
   123  		}
   124  
   125  		if len(res) == 0 {
   126  			break
   127  		}
   128  	}
   129  
   130  	return totalAcquire
   131  }
   132  
   133  // async janitor goroutine
   134  func janitor(c boskosClient, buffer chan string, fn clean) {
   135  	for {
   136  		proj := <-buffer
   137  
   138  		dest := common.Free
   139  		if err := fn(proj); err != nil {
   140  			logrus.WithError(err).Error("janitor.py failed!")
   141  			dest = common.Dirty
   142  		}
   143  
   144  		if err := c.ReleaseOne(proj, dest); err != nil {
   145  			logrus.WithError(err).Error("boskos release failed!")
   146  		}
   147  	}
   148  }