common/tsmon/target/flags.go

// Copyright 2015 The LUCI Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package target

import (
	"flag"
	"os"
	"regexp"
	"strings"

	"go.chromium.org/luci/common/tsmon/types"
)

// SysInfo overrides system's hostname and region for tests.
type SysInfo struct {
	Hostname string
	Region   string
}

// Flags defines command line flags related to tsmon targets.  Use NewFlags()
// to get a Flags struct with sensible default values.
type Flags struct {
	TargetType      types.TargetType
	DeviceHostname  string
	DeviceRegion    string
	DeviceRole      string
	DeviceNetwork   string
	TaskServiceName string
	TaskJobName     string
	TaskRegion      string
	TaskHostname    string
	TaskNumber      int
	AutoGenHostname bool

	// If nil, system info is computed from the actual host. Used
	// in tests.
	SysInfo *SysInfo
}

// NewFlags returns a Flags struct with sensible default values.  Hostname,
// region and network flags are expensive to compute, so get assigned default
// values later in SetDefaultsFromHostname.
func NewFlags() Flags {
	return Flags{
		TargetType:      DeviceType,
		DeviceHostname:  "",
		DeviceRegion:    "",
		DeviceRole:      "default",
		DeviceNetwork:   "",
		TaskServiceName: "",
		TaskJobName:     "",
		TaskRegion:      "",
		TaskHostname:    "",
		TaskNumber:      0,
		AutoGenHostname: false,
	}
}

// SetDefaultsFromHostname computes the expensive default values for hostname,
// region and network fields.
func (fl *Flags) SetDefaultsFromHostname() {
	if fl.SysInfo == nil {
		hostname, region := getFQDN()
		fl.SysInfo = &SysInfo{Hostname: hostname, Region: region}
	}
	network := getNetwork(fl.SysInfo.Hostname)
	hostname := fl.SysInfo.Hostname

	if fl.DeviceHostname == "" {
		fl.DeviceHostname = hostname
	}
	if fl.DeviceRegion == "" {
		fl.DeviceRegion = fl.SysInfo.Region
	}
	if fl.DeviceNetwork == "" {
		fl.DeviceNetwork = network
	}
	if fl.TaskRegion == "" {
		fl.TaskRegion = fl.SysInfo.Region
	}
	if fl.TaskHostname == "" {
		fl.TaskHostname = hostname
	}
	if fl.AutoGenHostname {
		fl.DeviceHostname = "autogen:" + fl.DeviceHostname
		fl.TaskHostname = "autogen:" + fl.TaskHostname
	}
}

// Register adds tsmon target related flags to a FlagSet.
func (fl *Flags) Register(f *flag.FlagSet) {
	f.Var(&targetTypeFlag{fl}, "ts-mon-target-type",
		"the type of target that is being monitored ("+targetTypeEnum.Choices()+")")
	f.StringVar(&fl.DeviceHostname, "ts-mon-device-hostname", fl.DeviceHostname,
		"name of this device")
	f.StringVar(&fl.DeviceRegion, "ts-mon-device-region", fl.DeviceRegion,
		"name of the region this devices lives in")
	f.StringVar(&fl.DeviceRole, "ts-mon-device-role", fl.DeviceRole,
		"role of the device")
	f.StringVar(&fl.DeviceNetwork, "ts-mon-device-network", fl.DeviceNetwork,
		"name of the network this device is connected to")
	f.StringVar(&fl.TaskServiceName, "ts-mon-task-service-name", fl.TaskServiceName,
		"name of the service being monitored")
	f.StringVar(&fl.TaskJobName, "ts-mon-task-job-name", fl.TaskJobName,
		"name of this job instance of the task")
	f.StringVar(&fl.TaskRegion, "ts-mon-task-region", fl.TaskRegion,
		"name of the region in which this task is running")
	f.StringVar(&fl.TaskHostname, "ts-mon-task-hostname", fl.TaskHostname,
		"name of the host on which this task is running")
	f.IntVar(&fl.TaskNumber, "ts-mon-task-number", fl.TaskNumber,
		"number (e.g. for replication) of this instance of this task")
	f.BoolVar(&fl.AutoGenHostname, "ts-mon-autogen-hostname", fl.AutoGenHostname,
		"Indicate that the hostname is autogenerated. "+
			"This option must be set on autoscaled GCE VMs, Kubernetes pods, "+
			"or any other hosts with dynamically generated names.")
}

func getFQDN() (string, string) {
	if hostname, err := os.Hostname(); err == nil {
		parts := strings.Split(hostname, ".")
		if len(parts) > 1 {
			return strings.ToLower(parts[0]), strings.ToLower(parts[1])
		}
		return strings.ToLower(hostname), "unknown"
	}
	return "unknown", "unknown"
}

func getNetwork(hostname string) string {
	// TODO(vadimsh): Move this to "hardcoded/chromeinfra" package.
	if m := regexp.MustCompile(`^([\w-]*?-[acm]|master)(\d+)a?$`).FindStringSubmatch(hostname); m != nil {
		return m[2]
	}
	return ""
}