github.com/shashidharatd/test-infra@v0.0.0-20171006011030-71304e1ca560/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  	poolSize       = 100 // Maximum concurrent janitor goroutines TODO(krzyzacy): should remove this limit
    32  	bufferSize     = 1   // Maximum holding resources
    33  	serviceAccount = flag.String("service-account", "", "Path to projects service account")
    34  )
    35  
    36  var rTypes common.ResTypes
    37  
    38  func init() {
    39  	flag.Var(&rTypes, "resource-type", "comma-separated list of resources need to be cleaned up")
    40  }
    41  
    42  func main() {
    43  	logrus.SetFormatter(&logrus.JSONFormatter{})
    44  	boskos := client.NewClient("Janitor", "http://boskos")
    45  	logrus.Info("Initialized boskos client!")
    46  
    47  	// Activate service account
    48  	flag.Parse()
    49  	if *serviceAccount == "" {
    50  		logrus.Fatal("--service-account cannot be empty!")
    51  	}
    52  
    53  	if len(rTypes) == 0 {
    54  		logrus.Fatal("--resource-type must not be empty!")
    55  	}
    56  
    57  	cmd := exec.Command("gcloud", "auth", "activate-service-account", "--key-file="+*serviceAccount)
    58  	if b, err := cmd.CombinedOutput(); err != nil {
    59  		logrus.WithError(err).Fatalf("fail to activate service account from %s :%s", *serviceAccount, string(b))
    60  	}
    61  
    62  	buffer := setup(boskos, poolSize, bufferSize, janitorClean)
    63  
    64  	for {
    65  		run(boskos, buffer, rTypes)
    66  		time.Sleep(time.Minute)
    67  	}
    68  }
    69  
    70  type clean func(string) error
    71  
    72  // Clean by janitor script
    73  func janitorClean(proj string) error {
    74  	cmd := exec.Command("/janitor.py", fmt.Sprintf("--project=%s", proj), "--hour=0")
    75  	b, err := cmd.CombinedOutput()
    76  	if err != nil {
    77  		logrus.Infof("janitor.py has some issue: %s", string(b))
    78  	}
    79  	return err
    80  }
    81  
    82  type boskosClient interface {
    83  	Acquire(rtype string, state string, dest string) (string, error)
    84  	ReleaseOne(name string, dest string) error
    85  }
    86  
    87  func setup(c boskosClient, janitorCount int, bufferSize int, cleanFunc clean) chan string {
    88  	buffer := make(chan string, 1)
    89  	for i := 0; i < janitorCount; i++ {
    90  		go janitor(c, buffer, cleanFunc)
    91  	}
    92  	return buffer
    93  }
    94  
    95  func run(c boskosClient, buffer chan string, rtypes []string) int {
    96  	totalAcquire := 0
    97  	res := make(map[string]int)
    98  	for _, s := range rtypes {
    99  		res[s] = 0
   100  	}
   101  
   102  	for {
   103  		for r := range res {
   104  			if proj, err := c.Acquire(r, "dirty", "cleaning"); err != nil {
   105  				logrus.WithError(err).Error("boskos acquire failed!")
   106  				totalAcquire += res[r]
   107  				delete(res, r)
   108  			} else if proj == "" {
   109  				totalAcquire += res[r]
   110  				delete(res, r)
   111  			} else {
   112  				buffer <- proj // will block until buffer has a free slot
   113  				res[r]++
   114  			}
   115  		}
   116  
   117  		if len(res) == 0 {
   118  			break
   119  		}
   120  	}
   121  
   122  	return totalAcquire
   123  }
   124  
   125  // async janitor goroutine
   126  func janitor(c boskosClient, buffer chan string, fn clean) {
   127  	for {
   128  		proj := <-buffer
   129  
   130  		dest := "free"
   131  		if err := fn(proj); err != nil {
   132  			logrus.WithError(err).Error("janitor.py failed!")
   133  			dest = "dirty"
   134  		}
   135  
   136  		if err := c.ReleaseOne(proj, dest); err != nil {
   137  			logrus.WithError(err).Error("boskos release failed!")
   138  		}
   139  	}
   140  }