vitess.io/vitess@v0.16.2/go/vt/vttest/environment.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 vttest
    18  
    19  import (
    20  	"fmt"
    21  	"math/rand"
    22  	"os"
    23  	"path"
    24  	"strings"
    25  	"time"
    26  
    27  	"vitess.io/vitess/go/vt/proto/vttest"
    28  
    29  	// we use gRPC everywhere, so import the vtgate client.
    30  	_ "vitess.io/vitess/go/vt/vtgate/grpcvtgateconn"
    31  )
    32  
    33  // Environment is the interface that customizes the global settings for
    34  // the test cluster. Usually the same environment settings are shared by
    35  // all the LocalCluster instances in a given test suite, with each instance
    36  // receiving a different Config for specific tests.
    37  // For Environments that create temporary data on-disk and clean it up on
    38  // termination, a brand new instance of Environment should be passed to
    39  // each LocalCluster.
    40  type Environment interface {
    41  	// BinaryPath returns the full path to the given executable
    42  	BinaryPath(bin string) string
    43  
    44  	// MySQLManager is the constructor for the MySQL manager that will
    45  	// be used by the cluster. The manager must take care of initializing
    46  	// and destructing the MySQL instance(s) that will be used by the cluster.
    47  	// See: vttest.MySQLManager for the interface the manager must implement
    48  	MySQLManager(mycnf []string, snapshot string) (MySQLManager, error)
    49  
    50  	// TopoManager is the constructor for the Topology manager that will
    51  	// be used by the cluster. It's only used when we run the local cluster with
    52  	// a remote topo server instead of in-memory topo server within vtcombo process
    53  	// See: vttest.TopoManager for the interface of topo manager
    54  	TopoManager(topoImplementation, topoServerAddress, topoRoot string, topology *vttest.VTTestTopology) TopoManager
    55  
    56  	// Directory is the path where the local cluster will store all its
    57  	// data and metadata. For local testing, this should probably be an
    58  	// unique temporary directory.
    59  	Directory() string
    60  
    61  	// LogDirectory is the directory where logs for all services in the
    62  	// cluster will be stored.
    63  	LogDirectory() string
    64  
    65  	// VtcomoboArguments are the extra commandline arguments that will be
    66  	// passed to `vtcombo`
    67  	VtcomboArguments() []string
    68  
    69  	// ProcessHealthCheck returns a HealthChecker for the given service.
    70  	// The HealthChecker takes an address and attempts to check whether
    71  	// the service is up and healthy.
    72  	// If a given service does not require any custom health checks,
    73  	// nil can be returned.
    74  	ProcessHealthCheck(name string) HealthChecker
    75  
    76  	// DefaultProtocol is the protocol used to communicate with the
    77  	// Vitess cluster. This is usually "grpc".
    78  	DefaultProtocol() string
    79  
    80  	// PortForProtocol returns the listening port for a given service
    81  	// on the given protocol. If protocol is empty, the default protocol
    82  	// for each service is assumed.
    83  	PortForProtocol(name, protocol string) int
    84  
    85  	// EnvVars returns the environment variables that will be passed
    86  	// to all Vitess processes spawned by the local cluster. These variables
    87  	// always take precedence over the variables inherited from the current
    88  	// process.
    89  	EnvVars() []string
    90  
    91  	// TearDown is called during LocalCluster.TearDown() to cleanup
    92  	// any temporary data in the environment. Environments that can
    93  	// last through several test runs do not need to implement it.
    94  	TearDown() error
    95  }
    96  
    97  // LocalTestEnv is an Environment implementation for local testing
    98  // See: NewLocalTestEnv()
    99  type LocalTestEnv struct {
   100  	BasePort     int
   101  	TmpPath      string
   102  	DefaultMyCnf []string
   103  	Env          []string
   104  }
   105  
   106  // DefaultMySQLFlavor is the MySQL flavor used by vttest when MYSQL_FLAVOR is not
   107  // set in the environment
   108  const DefaultMySQLFlavor = "MySQL56"
   109  
   110  // GetMySQLOptions returns the default option set for the given MySQL
   111  // flavor. If flavor is not set, the value from the `MYSQL_FLAVOR` env
   112  // variable is used, and if this is not set, DefaultMySQLFlavor will
   113  // be used.
   114  // Returns the name of the MySQL flavor being used, the set of MySQL CNF
   115  // files specific to this flavor, and any errors.
   116  func GetMySQLOptions(flavor string) (string, []string, error) {
   117  	if flavor == "" {
   118  		flavor = os.Getenv("MYSQL_FLAVOR")
   119  	}
   120  
   121  	if flavor == "" {
   122  		flavor = DefaultMySQLFlavor
   123  	}
   124  
   125  	mycnf := []string{}
   126  	mycnf = append(mycnf, "config/mycnf/test-suite.cnf")
   127  
   128  	for i, cnf := range mycnf {
   129  		mycnf[i] = path.Join(os.Getenv("VTROOT"), cnf)
   130  	}
   131  
   132  	return flavor, mycnf, nil
   133  }
   134  
   135  // EnvVars implements EnvVars for LocalTestEnv
   136  func (env *LocalTestEnv) EnvVars() []string {
   137  	return env.Env
   138  }
   139  
   140  // BinaryPath implements BinaryPath for LocalTestEnv
   141  func (env *LocalTestEnv) BinaryPath(binary string) string {
   142  	return path.Join(os.Getenv("VTROOT"), "bin", binary)
   143  }
   144  
   145  // MySQLManager implements MySQLManager for LocalTestEnv
   146  func (env *LocalTestEnv) MySQLManager(mycnf []string, snapshot string) (MySQLManager, error) {
   147  	return &Mysqlctl{
   148  		Binary:    env.BinaryPath("mysqlctl"),
   149  		InitFile:  path.Join(os.Getenv("VTROOT"), "config/init_db.sql"),
   150  		Directory: env.TmpPath,
   151  		Port:      env.PortForProtocol("mysql", ""),
   152  		MyCnf:     append(env.DefaultMyCnf, mycnf...),
   153  		Env:       env.EnvVars(),
   154  		UID:       1,
   155  	}, nil
   156  }
   157  
   158  // TopoManager implements TopoManager for LocalTestEnv
   159  func (env *LocalTestEnv) TopoManager(topoImplementation, topoServerAddress, topoRoot string, topology *vttest.VTTestTopology) TopoManager {
   160  	return &Topoctl{
   161  		TopoImplementation:      topoImplementation,
   162  		TopoGlobalServerAddress: topoServerAddress,
   163  		TopoGlobalRoot:          topoRoot,
   164  		Topology:                topology,
   165  	}
   166  }
   167  
   168  // DefaultProtocol implements DefaultProtocol for LocalTestEnv.
   169  // It is always GRPC.
   170  func (env *LocalTestEnv) DefaultProtocol() string {
   171  	return "grpc"
   172  }
   173  
   174  // PortForProtocol implements PortForProtocol for LocalTestEnv.
   175  func (env *LocalTestEnv) PortForProtocol(name, proto string) int {
   176  	switch name {
   177  	case "vtcombo":
   178  		if proto == "grpc" {
   179  			return env.BasePort + 1
   180  		}
   181  		return env.BasePort
   182  
   183  	case "mysql":
   184  		return env.BasePort + 2
   185  
   186  	case "vtcombo_mysql_port":
   187  		return env.BasePort + 3
   188  
   189  	default:
   190  		panic("unknown service name: " + name)
   191  	}
   192  }
   193  
   194  // ProcessHealthCheck implements ProcessHealthCheck for LocalTestEnv.
   195  // By default, it performs no service-specific health checks
   196  func (env *LocalTestEnv) ProcessHealthCheck(name string) HealthChecker {
   197  	return nil
   198  }
   199  
   200  // VtcomboArguments implements VtcomboArguments for LocalTestEnv.
   201  func (env *LocalTestEnv) VtcomboArguments() []string {
   202  	return []string{
   203  		"--service_map", strings.Join(
   204  			[]string{"grpc-vtgateservice", "grpc-vtctl", "grpc-vtctld"}, ",",
   205  		),
   206  	}
   207  }
   208  
   209  // LogDirectory implements LogDirectory for LocalTestEnv.
   210  func (env *LocalTestEnv) LogDirectory() string {
   211  	return path.Join(env.TmpPath, "logs")
   212  }
   213  
   214  // Directory implements Directory for LocalTestEnv.
   215  func (env *LocalTestEnv) Directory() string {
   216  	return env.TmpPath
   217  }
   218  
   219  // TearDown implements TearDown for LocalTestEnv
   220  func (env *LocalTestEnv) TearDown() error {
   221  	return os.RemoveAll(env.TmpPath)
   222  }
   223  
   224  func tmpdir(dataroot string) (dir string, err error) {
   225  	dir, err = os.MkdirTemp(dataroot, "vttest")
   226  	return
   227  }
   228  
   229  func randomPort() int {
   230  	v := rand.Int31n(20000)
   231  	return int(v + 10000)
   232  }
   233  
   234  // NewLocalTestEnv returns an instance of the default test environment used
   235  // for local testing Vitess. The defaults are as follows:
   236  // - Directory() is a random temporary directory in VTDATAROOT, which is cleaned
   237  // up when closing the Environment.
   238  // - LogDirectory() is the `logs` subdir inside Directory()
   239  // - The MySQL flavor is set to `flavor`. If the argument is not set, it will
   240  // default to the value of MYSQL_FLAVOR, and if this variable is not set, to
   241  // DefaultMySQLFlavor
   242  // - PortForProtocol() will return ports based off the given basePort. If basePort
   243  // is zero, a random port between 10000 and 20000 will be chosen.
   244  // - DefaultProtocol() is always "grpc"
   245  // - ProcessHealthCheck() performs no service-specific health checks
   246  // - BinaryPath() will look up the default Vitess binaries in VTROOT
   247  // - MySQLManager() will return a vttest.Mysqlctl instance, configured with the
   248  // given MySQL flavor. This will use the `mysqlctl` command to initialize and
   249  // teardown a single mysqld instance.
   250  func NewLocalTestEnv(flavor string, basePort int) (*LocalTestEnv, error) {
   251  	directory, err := tmpdir(os.Getenv("VTDATAROOT"))
   252  	if err != nil {
   253  		return nil, err
   254  	}
   255  	return NewLocalTestEnvWithDirectory(flavor, basePort, directory)
   256  }
   257  
   258  // NewLocalTestEnvWithDirectory returns a new instance of the default test
   259  // environment with a directory explicitly specified.
   260  func NewLocalTestEnvWithDirectory(flavor string, basePort int, directory string) (*LocalTestEnv, error) {
   261  	if _, err := os.Stat(path.Join(directory, "logs")); os.IsNotExist(err) {
   262  		err := os.Mkdir(path.Join(directory, "logs"), 0700)
   263  		if err != nil {
   264  			return nil, err
   265  		}
   266  	}
   267  
   268  	flavor, mycnf, err := GetMySQLOptions(flavor)
   269  	if err != nil {
   270  		return nil, err
   271  	}
   272  
   273  	if basePort == 0 {
   274  		basePort = randomPort()
   275  	}
   276  
   277  	return &LocalTestEnv{
   278  		BasePort:     basePort,
   279  		TmpPath:      directory,
   280  		DefaultMyCnf: mycnf,
   281  		Env: []string{
   282  			fmt.Sprintf("VTDATAROOT=%s", directory),
   283  			fmt.Sprintf("MYSQL_FLAVOR=%s", flavor),
   284  		},
   285  	}, nil
   286  }
   287  
   288  func defaultEnvFactory() (Environment, error) {
   289  	return NewLocalTestEnv("", 0)
   290  }
   291  
   292  func init() {
   293  	rand.Seed(time.Now().UnixNano())
   294  }
   295  
   296  // NewDefaultEnv is an user-configurable callback that returns a new Environment
   297  // instance with the default settings.
   298  // This callback is only used in cases where the user hasn't explicitly set
   299  // the Env variable when initializing a LocalCluster
   300  var NewDefaultEnv = defaultEnvFactory