github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/cmd/reduce/main.go (about)

     1  // Copyright 2019 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  // reduce reduces SQL passed over stdin using cockroach demo. The input is
    12  // simplified such that the contains argument is present as an error during SQL
    13  // execution.
    14  package main
    15  
    16  import (
    17  	"context"
    18  	"flag"
    19  	"fmt"
    20  	"io"
    21  	"io/ioutil"
    22  	"log"
    23  	"os"
    24  	"os/exec"
    25  	"regexp"
    26  	"runtime"
    27  	"strings"
    28  	"time"
    29  
    30  	"github.com/cockroachdb/cockroach/pkg/testutils/reduce"
    31  	"github.com/cockroachdb/cockroach/pkg/testutils/reduce/reducesql"
    32  	"github.com/cockroachdb/errors"
    33  )
    34  
    35  var (
    36  	// Some quick benchmarks show that somewhere around 1/3 of NumCPUs
    37  	// performs best. This can probably be tweaked with benchmarks from
    38  	// other machines, but is probably a good place to start.
    39  	goroutines = func() int {
    40  		// Round up by adding 2.
    41  		// Num CPUs -> n:
    42  		// 1-3: 1
    43  		// 4-6: 2
    44  		// 7-9: 3
    45  		// etc.
    46  		n := (runtime.NumCPU() + 2) / 3
    47  		if n < 1 {
    48  			n = 1
    49  		}
    50  		return n
    51  	}()
    52  	flags    = flag.NewFlagSet(os.Args[0], flag.ExitOnError)
    53  	path     = flags.String("path", "./cockroach", "path to cockroach binary")
    54  	verbose  = flags.Bool("v", false, "log progress")
    55  	contains = flags.String("contains", "", "error regex to search for")
    56  	unknown  = flags.Bool("unknown", false, "print unknown types during walk")
    57  	workers  = flags.Int("goroutines", goroutines, "number of worker goroutines (defaults to NumCPU/3")
    58  )
    59  
    60  func usage() {
    61  	fmt.Fprintf(flags.Output(), "Usage of %s:\n", os.Args[0])
    62  	flags.PrintDefaults()
    63  	os.Exit(1)
    64  }
    65  
    66  func main() {
    67  	if err := flags.Parse(os.Args[1:]); err != nil {
    68  		usage()
    69  	}
    70  	if *contains == "" {
    71  		fmt.Print("missing contains\n\n")
    72  		usage()
    73  	}
    74  	reducesql.LogUnknown = *unknown
    75  	out, err := reduceSQL(*path, *contains, *workers, *verbose)
    76  	if err != nil {
    77  		log.Fatal(err)
    78  	}
    79  	fmt.Println(out)
    80  }
    81  
    82  func reduceSQL(path, contains string, workers int, verbose bool) (string, error) {
    83  	containsRE, err := regexp.Compile(contains)
    84  	if err != nil {
    85  		return "", err
    86  	}
    87  	var input []byte
    88  	{
    89  		done := make(chan struct{}, 1)
    90  		go func() {
    91  			select {
    92  			case <-done:
    93  			case <-time.After(5 * time.Second):
    94  				log.Fatal("timeout waiting for input on stdin")
    95  			}
    96  		}()
    97  		input, err = ioutil.ReadAll(os.Stdin)
    98  		done <- struct{}{}
    99  		if err != nil {
   100  			return "", err
   101  		}
   102  	}
   103  
   104  	// Pretty print the input so the file size comparison is useful.
   105  	inputSQL, err := reducesql.Pretty(input)
   106  	if err != nil {
   107  		return "", err
   108  	}
   109  
   110  	var logger io.Writer
   111  	if verbose {
   112  		logger = os.Stderr
   113  		fmt.Fprintf(logger, "input SQL pretty printed, %d bytes -> %d bytes\n", len(input), len(inputSQL))
   114  	}
   115  
   116  	interesting := func(ctx context.Context, f reduce.File) bool {
   117  		// Disable telemetry and license generation.
   118  		cmd := exec.CommandContext(ctx, path,
   119  			"demo",
   120  			"--empty",
   121  			"--disable-demo-license",
   122  		)
   123  		cmd.Env = []string{"COCKROACH_SKIP_ENABLING_DIAGNOSTIC_REPORTING", "true"}
   124  		sql := string(f)
   125  		if !strings.HasSuffix(sql, ";") {
   126  			sql += ";"
   127  		}
   128  		cmd.Stdin = strings.NewReader(sql)
   129  		out, err := cmd.CombinedOutput()
   130  		switch {
   131  		case errors.HasType(err, (*exec.Error)(nil)):
   132  			if errors.Is(err, exec.ErrNotFound) {
   133  				log.Fatal(err)
   134  			}
   135  		case errors.HasType(err, (*os.PathError)(nil)):
   136  			log.Fatal(err)
   137  		}
   138  		return containsRE.Match(out)
   139  	}
   140  
   141  	out, err := reduce.Reduce(logger, reduce.File(inputSQL), interesting, workers, reduce.ModeInteresting, reducesql.SQLPasses...)
   142  	return string(out), err
   143  }