github.com/pingcap/chaos@v0.0.0-20190710112158-c86faf4b3719/cmd/util/suit.go (about)

     1  package util
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"log"
     7  	"os"
     8  	"os/signal"
     9  	"strings"
    10  	"syscall"
    11  
    12  	"github.com/pingcap/chaos/pkg/control"
    13  	"github.com/pingcap/chaos/pkg/core"
    14  	"github.com/pingcap/chaos/pkg/nemesis"
    15  	"github.com/pingcap/chaos/pkg/verify"
    16  )
    17  
    18  // Suit is a basic chaos testing suit with configurations to run chaos.
    19  type Suit struct {
    20  	*control.Config
    21  	core.ClientCreator
    22  	// nemesis, seperated by comma.
    23  	Nemesises string
    24  
    25  	VerifySuit verify.Suit
    26  }
    27  
    28  // Run runs the suit.
    29  func (suit *Suit) Run(ctx context.Context, nodes []string) {
    30  	var nemesisGens []core.NemesisGenerator
    31  	for _, name := range strings.Split(suit.Nemesises, ",") {
    32  		var g core.NemesisGenerator
    33  		name := strings.TrimSpace(name)
    34  		if len(name) == 0 {
    35  			continue
    36  		}
    37  
    38  		switch name {
    39  		case "random_kill", "all_kill", "minor_kill", "major_kill":
    40  			g = nemesis.NewKillGenerator(suit.Config.DB, name)
    41  		case "random_drop", "all_drop", "minor_drop", "major_drop":
    42  			g = nemesis.NewDropGenerator(name)
    43  		default:
    44  			log.Fatalf("invalid nemesis generator %s", name)
    45  		}
    46  
    47  		nemesisGens = append(nemesisGens, g)
    48  	}
    49  
    50  	sctx, cancel := context.WithCancel(ctx)
    51  
    52  	if len(nodes) != 0 {
    53  		suit.Config.Nodes = nodes
    54  	} else {
    55  		// By default, we run TiKV/TiDB cluster on 5 nodes.
    56  		for i := 1; i <= 5; i++ {
    57  			name := fmt.Sprintf("n%d", i)
    58  			suit.Config.Nodes = append(suit.Config.Nodes, name)
    59  		}
    60  	}
    61  
    62  	c := control.NewController(
    63  		sctx,
    64  		suit.Config,
    65  		suit.ClientCreator,
    66  		nemesisGens,
    67  		suit.VerifySuit,
    68  	)
    69  
    70  	sigs := make(chan os.Signal, 1)
    71  	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)
    72  
    73  	go func() {
    74  		<-sigs
    75  		c.Close()
    76  		cancel()
    77  	}()
    78  
    79  	c.Run()
    80  }