go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/common/tsmon/target/flags.go (about)

     1  // Copyright 2015 The LUCI Authors.
     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  
    15  package target
    16  
    17  import (
    18  	"flag"
    19  	"os"
    20  	"regexp"
    21  	"strings"
    22  
    23  	"go.chromium.org/luci/common/tsmon/types"
    24  )
    25  
    26  // SysInfo overrides system's hostname and region for tests.
    27  type SysInfo struct {
    28  	Hostname string
    29  	Region   string
    30  }
    31  
    32  // Flags defines command line flags related to tsmon targets.  Use NewFlags()
    33  // to get a Flags struct with sensible default values.
    34  type Flags struct {
    35  	TargetType      types.TargetType
    36  	DeviceHostname  string
    37  	DeviceRegion    string
    38  	DeviceRole      string
    39  	DeviceNetwork   string
    40  	TaskServiceName string
    41  	TaskJobName     string
    42  	TaskRegion      string
    43  	TaskHostname    string
    44  	TaskNumber      int
    45  	AutoGenHostname bool
    46  
    47  	// If nil, system info is computed from the actual host. Used
    48  	// in tests.
    49  	SysInfo *SysInfo
    50  }
    51  
    52  // NewFlags returns a Flags struct with sensible default values.  Hostname,
    53  // region and network flags are expensive to compute, so get assigned default
    54  // values later in SetDefaultsFromHostname.
    55  func NewFlags() Flags {
    56  	return Flags{
    57  		TargetType:      DeviceType,
    58  		DeviceHostname:  "",
    59  		DeviceRegion:    "",
    60  		DeviceRole:      "default",
    61  		DeviceNetwork:   "",
    62  		TaskServiceName: "",
    63  		TaskJobName:     "",
    64  		TaskRegion:      "",
    65  		TaskHostname:    "",
    66  		TaskNumber:      0,
    67  		AutoGenHostname: false,
    68  	}
    69  }
    70  
    71  // SetDefaultsFromHostname computes the expensive default values for hostname,
    72  // region and network fields.
    73  func (fl *Flags) SetDefaultsFromHostname() {
    74  	if fl.SysInfo == nil {
    75  		hostname, region := getFQDN()
    76  		fl.SysInfo = &SysInfo{Hostname: hostname, Region: region}
    77  	}
    78  	network := getNetwork(fl.SysInfo.Hostname)
    79  	hostname := fl.SysInfo.Hostname
    80  
    81  	if fl.DeviceHostname == "" {
    82  		fl.DeviceHostname = hostname
    83  	}
    84  	if fl.DeviceRegion == "" {
    85  		fl.DeviceRegion = fl.SysInfo.Region
    86  	}
    87  	if fl.DeviceNetwork == "" {
    88  		fl.DeviceNetwork = network
    89  	}
    90  	if fl.TaskRegion == "" {
    91  		fl.TaskRegion = fl.SysInfo.Region
    92  	}
    93  	if fl.TaskHostname == "" {
    94  		fl.TaskHostname = hostname
    95  	}
    96  	if fl.AutoGenHostname {
    97  		fl.DeviceHostname = "autogen:" + fl.DeviceHostname
    98  		fl.TaskHostname = "autogen:" + fl.TaskHostname
    99  	}
   100  }
   101  
   102  // Register adds tsmon target related flags to a FlagSet.
   103  func (fl *Flags) Register(f *flag.FlagSet) {
   104  	f.Var(&targetTypeFlag{fl}, "ts-mon-target-type",
   105  		"the type of target that is being monitored ("+targetTypeEnum.Choices()+")")
   106  	f.StringVar(&fl.DeviceHostname, "ts-mon-device-hostname", fl.DeviceHostname,
   107  		"name of this device")
   108  	f.StringVar(&fl.DeviceRegion, "ts-mon-device-region", fl.DeviceRegion,
   109  		"name of the region this devices lives in")
   110  	f.StringVar(&fl.DeviceRole, "ts-mon-device-role", fl.DeviceRole,
   111  		"role of the device")
   112  	f.StringVar(&fl.DeviceNetwork, "ts-mon-device-network", fl.DeviceNetwork,
   113  		"name of the network this device is connected to")
   114  	f.StringVar(&fl.TaskServiceName, "ts-mon-task-service-name", fl.TaskServiceName,
   115  		"name of the service being monitored")
   116  	f.StringVar(&fl.TaskJobName, "ts-mon-task-job-name", fl.TaskJobName,
   117  		"name of this job instance of the task")
   118  	f.StringVar(&fl.TaskRegion, "ts-mon-task-region", fl.TaskRegion,
   119  		"name of the region in which this task is running")
   120  	f.StringVar(&fl.TaskHostname, "ts-mon-task-hostname", fl.TaskHostname,
   121  		"name of the host on which this task is running")
   122  	f.IntVar(&fl.TaskNumber, "ts-mon-task-number", fl.TaskNumber,
   123  		"number (e.g. for replication) of this instance of this task")
   124  	f.BoolVar(&fl.AutoGenHostname, "ts-mon-autogen-hostname", fl.AutoGenHostname,
   125  		"Indicate that the hostname is autogenerated. "+
   126  			"This option must be set on autoscaled GCE VMs, Kubernetes pods, "+
   127  			"or any other hosts with dynamically generated names.")
   128  }
   129  
   130  func getFQDN() (string, string) {
   131  	if hostname, err := os.Hostname(); err == nil {
   132  		parts := strings.Split(hostname, ".")
   133  		if len(parts) > 1 {
   134  			return strings.ToLower(parts[0]), strings.ToLower(parts[1])
   135  		}
   136  		return strings.ToLower(hostname), "unknown"
   137  	}
   138  	return "unknown", "unknown"
   139  }
   140  
   141  func getNetwork(hostname string) string {
   142  	// TODO(vadimsh): Move this to "hardcoded/chromeinfra" package.
   143  	if m := regexp.MustCompile(`^([\w-]*?-[acm]|master)(\d+)a?$`).FindStringSubmatch(hostname); m != nil {
   144  		return m[2]
   145  	}
   146  	return ""
   147  }