vitess.io/vitess@v0.16.2/go/vt/servenv/service_map.go (about)

     1  /*
     2  Copyright 2019 The Vitess 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 servenv
    18  
    19  import (
    20  	"github.com/spf13/pflag"
    21  
    22  	"vitess.io/vitess/go/vt/log"
    23  )
    24  
    25  var (
    26  	serviceMapFlag []string
    27  
    28  	// serviceMap is the used version of the service map.
    29  	// init() functions can add default values to it (using InitServiceMap).
    30  	// service_map command line parameter will alter the map.
    31  	// Can only be used after servenv.Init has been called.
    32  	serviceMap = make(map[string]bool)
    33  )
    34  
    35  // RegisterServiceMapFlag registers an OnParse hook to install the
    36  // `--service_map` flag for a given cmd. It must be called before ParseFlags or
    37  // ParseFlagsWithArgs.
    38  func RegisterServiceMapFlag() {
    39  	OnParse(func(fs *pflag.FlagSet) {
    40  		fs.StringSliceVar(&serviceMapFlag, "service_map", serviceMapFlag, "comma separated list of services to enable (or disable if prefixed with '-') Example: grpc-queryservice")
    41  	})
    42  	OnInit(updateServiceMap)
    43  }
    44  
    45  // InitServiceMap will set the default value for a protocol/name to be served.
    46  func InitServiceMap(protocol, name string) {
    47  	serviceMap[protocol+"-"+name] = true
    48  }
    49  
    50  // updateServiceMap takes the command line parameter, and updates the
    51  // ServiceMap accordingly
    52  func updateServiceMap() {
    53  	for _, s := range serviceMapFlag {
    54  		if s[0] == '-' {
    55  			delete(serviceMap, s[1:])
    56  		} else {
    57  			serviceMap[s] = true
    58  		}
    59  	}
    60  }
    61  
    62  // checkServiceMap returns if we should register a RPC service
    63  // (and also logs how to enable / disable it)
    64  func checkServiceMap(protocol, name string) bool {
    65  	if serviceMap[protocol+"-"+name] {
    66  		log.Infof("Registering %v for %v, disable it with -%v-%v service_map parameter", name, protocol, protocol, name)
    67  		return true
    68  	}
    69  	log.Infof("Not registering %v for %v, enable it with %v-%v service_map parameter", name, protocol, protocol, name)
    70  	return false
    71  }