<commit_before>// Copyright 2013 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// +build !windows,!plan9

package interp_test

import (
	"bytes"
	"fmt"
	"go/build"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"golang.org/x/tools/go/loader"
	"golang.org/x/tools/go/ssa"
	"golang.org/x/tools/go/ssa/interp"
	"golang.org/x/tools/go/ssa/ssautil"
	"golang.org/x/tools/go/types"
)

// Each line contains a space-separated list of $GOROOT/test/
// filenames comprising the main package of a program.
// They are ordered quickest-first, roughly.
//
// TODO(adonovan): integrate into the $GOROOT/test driver scripts,
// golden file checking, etc.
var gorootTestTests = []string{
	"235.go",
	"alias1.go",
	"chancap.go",
	"func5.go",
	"func6.go",
	"func7.go",
	"func8.go",
	"helloworld.go",
	"varinit.go",
	"escape3.go",
	"initcomma.go",
	"cmp.go",
	"compos.go",
	"turing.go",
	"indirect.go",
	"complit.go",
	"for.go",
	"struct0.go",
	"intcvt.go",
	"printbig.go",
	"deferprint.go",
	"escape.go",
	"range.go",
	"const4.go",
	"float_lit.go",
	"bigalg.go",
	"decl.go",
	"if.go",
	"named.go",
	"bigmap.go",
	"func.go",
	"reorder2.go",
	"closure.go",
	"gc.go",
	"simassign.go",
	"iota.go",
	"nilptr2.go",
	"goprint.go", // doesn't actually assert anything (cmpout)
	"utf.go",
	"method.go",
	"char_lit.go",
	"env.go",
	"int_lit.go",
	"string_lit.go",
	"defer.go",
	"typeswitch.go",
	"stringrange.go",
	"reorder.go",
	"method3.go",
	"literal.go",
	"nul1.go", // doesn't actually assert anything (errorcheckoutput)
	"zerodivide.go",
	"convert.go",
	"convT2X.go",
	"switch.go",
	"initialize.go",
	"ddd.go",
	"blank.go", // partly disabled
	"map.go",
	"closedchan.go",
	"divide.go",
	"rename.go",
	"const3.go",
	"nil.go",
	"recover.go", // reflection parts disabled
	"recover1.go",
	"recover2.go",
	"recover3.go",
	"typeswitch1.go",
	"floatcmp.go",
	"crlf.go", // doesn't actually assert anything (runoutput)
	// Slow tests follow.
	"bom.go", // ~1.7s
	"gc1.go", // ~1.7s
	"cmplxdivide.go cmplxdivide1.go", // ~2.4s

	// Working, but not worth enabling:
	// "append.go",    // works, but slow (15s).
	// "gc2.go",       // works, but slow, and cheats on the memory check.
	// "sigchld.go",   // works, but only on POSIX.
	// "peano.go",     // works only up to n=9, and slow even then.
	// "stack.go",     // works, but too slow (~30s) by default.
	// "solitaire.go", // works, but too slow (~30s).
	// "const.go",     // works but for but one bug: constant folder doesn't consider representations.
	// "init1.go",     // too slow (80s) and not that interesting. Cheats on ReadMemStats check too.
	// "rotate.go rotate0.go", // emits source for a test
	// "rotate.go rotate1.go", // emits source for a test
	// "rotate.go rotate2.go", // emits source for a test
	// "rotate.go rotate3.go", // emits source for a test
	// "64bit.go",             // emits source for a test
	// "run.go",               // test driver, not a test.

	// Broken.  TODO(adonovan): fix.
	// copy.go         // very slow; but with N=4 quickly crashes, slice index out of range.
	// nilptr.go       // interp: V > uintptr not implemented. Slow test, lots of mem
	// args.go         // works, but requires specific os.Args from the driver.
	// index.go        // a template, not a real test.
	// mallocfin.go    // SetFinalizer not implemented.

	// TODO(adonovan): add tests from $GOROOT/test/* subtrees:
	// bench chan bugs fixedbugs interface ken.
}

// These are files in go.tools/go/ssa/interp/testdata/.
var testdataTests = []string{
	"boundmeth.go",
	"complit.go",
	"coverage.go",
	"defer.go",
	"fieldprom.go",
	"ifaceconv.go",
	"ifaceprom.go",
	"initorder.go",
	"methprom.go",
	"mrvchain.go",
	"range.go",
	"recover.go",
	"reflect.go",
	"static.go",
	"callstack.go",
}

// These are files and packages in $GOROOT/src/.
var gorootSrcTests = []string{
	"encoding/ascii85",
	"encoding/csv",
	"encoding/hex",
	"encoding/pem",
	"hash/crc32",
	// "testing", // TODO(adonovan): implement runtime.Goexit correctly
	"text/scanner",
	"unicode",

	// Too slow:
	// "container/ring",
	// "hash/adler32",

	// TODO(adonovan): packages with Examples require os.Pipe (unimplemented):
	// "unicode/utf8",
	// "log",
	// "path",
	// "flag",
}

type successPredicate func(exitcode int, output string) error

func run(t *testing.T, dir, input string, success successPredicate) bool {
	fmt.Printf("Input: %s\n", input)

	start := time.Now()

	var inputs []string
	for _, i := range strings.Split(input, " ") {
		if strings.HasSuffix(i, ".go") {
			i = dir + i
		}
		inputs = append(inputs, i)
	}

	var conf loader.Config
	if _, err := conf.FromArgs(inputs, true); err != nil {
		t.Errorf("FromArgs(%s) failed: %s", inputs, err)
		return false
	}

	conf.Import("runtime")

	// Print a helpful hint if we don't make it to the end.
	var hint string
	defer func() {
		if hint != "" {
			fmt.Println("FAIL")
			fmt.Println(hint)
		} else {
			fmt.Println("PASS")
		}

		interp.CapturedOutput = nil
	}()

	hint = fmt.Sprintf("To dump SSA representation, run:\n%% go build golang.org/x/tools/cmd/ssadump && ./ssadump -build=CFP %s\n", input)

	iprog, err := conf.Load()
	if err != nil {
		t.Errorf("conf.Load(%s) failed: %s", inputs, err)
		return false
	}

	prog := ssautil.CreateProgram(iprog, ssa.SanityCheckFunctions)
	prog.BuildAll()

	var mainPkg *ssa.Package
	var initialPkgs []*ssa.Package
	for _, info := range iprog.InitialPackages() {
		if info.Pkg.Path() == "runtime" {
			continue // not an initial package
		}
		p := prog.Package(info.Pkg)
		initialPkgs = append(initialPkgs, p)
		if mainPkg == nil && p.Func("main") != nil {
			mainPkg = p
		}
	}
	if mainPkg == nil {
		testmainPkg := prog.CreateTestMainPackage(initialPkgs...)
		if testmainPkg == nil {
			t.Errorf("CreateTestMainPackage(%s) returned nil", mainPkg)
			return false
		}
		if testmainPkg.Func("main") == nil {
			t.Errorf("synthetic testmain package has no main")
			return false
		}
		mainPkg = testmainPkg
	}

	var out bytes.Buffer
	interp.CapturedOutput = &out

	hint = fmt.Sprintf("To trace execution, run:\n%% go build golang.org/x/tools/cmd/ssadump && ./ssadump -build=C -run --interp=T %s\n", input)
	exitCode := interp.Interpret(mainPkg, 0, &types.StdSizes{8, 8}, inputs[0], []string{})

	// The definition of success varies with each file.
	if err := success(exitCode, out.String()); err != nil {
		t.Errorf("interp.Interpret(%s) failed: %s", inputs, err)
		return false
	}

	hint = "" // call off the hounds

	if false {
		fmt.Println(input, time.Since(start)) // test profiling
	}

	return true
}

const slash = string(os.PathSeparator)

func printFailures(failures []string) {
	if failures != nil {
		fmt.Println("The following tests failed:")
		for _, f := range failures {
			fmt.Printf("\t%s\n", f)
		}
	}
}

func success(exitcode int, output string) error {
	if exitcode != 0 {
		return fmt.Errorf("exit code was %d", exitcode)
	}
	if strings.Contains(output, "BUG") {
		return fmt.Errorf("exited zero but output contained 'BUG'")
	}
	return nil
}

// TestTestdataFiles runs the interpreter on testdata/*.go.
func TestTestdataFiles(t *testing.T) {
	var failures []string
	for _, input := range testdataTests {
		if !run(t, "testdata"+slash, input, success) {
			failures = append(failures, input)
		}
	}
	printFailures(failures)
}

// TestGorootTest runs the interpreter on $GOROOT/test/*.go.
func TestGorootTest(t *testing.T) {
	if testing.Short() {
		return // too slow (~30s)
	}

	var failures []string

	for _, input := range gorootTestTests {
		if !run(t, filepath.Join(build.Default.GOROOT, "test")+slash, input, success) {
			failures = append(failures, input)
		}
	}
	for _, input := range gorootSrcTests {
		if !run(t, filepath.Join(build.Default.GOROOT, "src")+slash, input, success) {
			failures = append(failures, input)
		}
	}
	printFailures(failures)
}

// TestTestmainPackage runs the interpreter on a synthetic "testmain" package.
func TestTestmainPackage(t *testing.T) {
	success := func(exitcode int, output string) error {
		if exitcode == 0 {
			return fmt.Errorf("unexpected success")
		}
		if !strings.Contains(output, "FAIL: TestFoo") {
			return fmt.Errorf("missing failure log for TestFoo")
		}
		if !strings.Contains(output, "FAIL: TestBar") {
			return fmt.Errorf("missing failure log for TestBar")
		}
		// TODO(adonovan): test benchmarks too
		return nil
	}
	run(t, "testdata"+slash, "a_test.go", success)
}

// CreateTestMainPackage should return nil if there were no tests.
func TestNullTestmainPackage(t *testing.T) {
	var conf loader.Config
	conf.CreateFromFilenames("", "testdata/b_test.go")
	iprog, err := conf.Load()
	if err != nil {
		t.Fatalf("CreatePackages failed: %s", err)
	}
	prog := ssautil.CreateProgram(iprog, ssa.SanityCheckFunctions)
	mainPkg := prog.Package(iprog.Created[0].Pkg)
	if mainPkg.Func("main") != nil {
		t.Fatalf("unexpected main function")
	}
	if prog.CreateTestMainPackage(mainPkg) != nil {
		t.Fatalf("CreateTestMainPackage returned non-nil")
	}
}
<commit_msg>go/ssa/interp: remove hash/crc32 from test suite, since it uses Examples<commit_after>// Copyright 2013 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// +build !windows,!plan9

package interp_test

import (
	"bytes"
	"fmt"
	"go/build"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"golang.org/x/tools/go/loader"
	"golang.org/x/tools/go/ssa"
	"golang.org/x/tools/go/ssa/interp"
	"golang.org/x/tools/go/ssa/ssautil"
	"golang.org/x/tools/go/types"
)

// Each line contains a space-separated list of $GOROOT/test/
// filenames comprising the main package of a program.
// They are ordered quickest-first, roughly.
//
// TODO(adonovan): integrate into the $GOROOT/test driver scripts,
// golden file checking, etc.
var gorootTestTests = []string{
	"235.go",
	"alias1.go",
	"chancap.go",
	"func5.go",
	"func6.go",
	"func7.go",
	"func8.go",
	"helloworld.go",
	"varinit.go",
	"escape3.go",
	"initcomma.go",
	"cmp.go",
	"compos.go",
	"turing.go",
	"indirect.go",
	"complit.go",
	"for.go",
	"struct0.go",
	"intcvt.go",
	"printbig.go",
	"deferprint.go",
	"escape.go",
	"range.go",
	"const4.go",
	"float_lit.go",
	"bigalg.go",
	"decl.go",
	"if.go",
	"named.go",
	"bigmap.go",
	"func.go",
	"reorder2.go",
	"closure.go",
	"gc.go",
	"simassign.go",
	"iota.go",
	"nilptr2.go",
	"goprint.go", // doesn't actually assert anything (cmpout)
	"utf.go",
	"method.go",
	"char_lit.go",
	"env.go",
	"int_lit.go",
	"string_lit.go",
	"defer.go",
	"typeswitch.go",
	"stringrange.go",
	"reorder.go",
	"method3.go",
	"literal.go",
	"nul1.go", // doesn't actually assert anything (errorcheckoutput)
	"zerodivide.go",
	"convert.go",
	"convT2X.go",
	"switch.go",
	"initialize.go",
	"ddd.go",
	"blank.go", // partly disabled
	"map.go",
	"closedchan.go",
	"divide.go",
	"rename.go",
	"const3.go",
	"nil.go",
	"recover.go", // reflection parts disabled
	"recover1.go",
	"recover2.go",
	"recover3.go",
	"typeswitch1.go",
	"floatcmp.go",
	"crlf.go", // doesn't actually assert anything (runoutput)
	// Slow tests follow.
	"bom.go", // ~1.7s
	"gc1.go", // ~1.7s
	"cmplxdivide.go cmplxdivide1.go", // ~2.4s

	// Working, but not worth enabling:
	// "append.go",    // works, but slow (15s).
	// "gc2.go",       // works, but slow, and cheats on the memory check.
	// "sigchld.go",   // works, but only on POSIX.
	// "peano.go",     // works only up to n=9, and slow even then.
	// "stack.go",     // works, but too slow (~30s) by default.
	// "solitaire.go", // works, but too slow (~30s).
	// "const.go",     // works but for but one bug: constant folder doesn't consider representations.
	// "init1.go",     // too slow (80s) and not that interesting. Cheats on ReadMemStats check too.
	// "rotate.go rotate0.go", // emits source for a test
	// "rotate.go rotate1.go", // emits source for a test
	// "rotate.go rotate2.go", // emits source for a test
	// "rotate.go rotate3.go", // emits source for a test
	// "64bit.go",             // emits source for a test
	// "run.go",               // test driver, not a test.

	// Broken.  TODO(adonovan): fix.
	// copy.go         // very slow; but with N=4 quickly crashes, slice index out of range.
	// nilptr.go       // interp: V > uintptr not implemented. Slow test, lots of mem
	// args.go         // works, but requires specific os.Args from the driver.
	// index.go        // a template, not a real test.
	// mallocfin.go    // SetFinalizer not implemented.

	// TODO(adonovan): add tests from $GOROOT/test/* subtrees:
	// bench chan bugs fixedbugs interface ken.
}

// These are files in go.tools/go/ssa/interp/testdata/.
var testdataTests = []string{
	"boundmeth.go",
	"complit.go",
	"coverage.go",
	"defer.go",
	"fieldprom.go",
	"ifaceconv.go",
	"ifaceprom.go",
	"initorder.go",
	"methprom.go",
	"mrvchain.go",
	"range.go",
	"recover.go",
	"reflect.go",
	"static.go",
	"callstack.go",
}

// These are files and packages in $GOROOT/src/.
var gorootSrcTests = []string{
	"encoding/ascii85",
	"encoding/csv",
	"encoding/hex",
	"encoding/pem",
	// "testing", // TODO(adonovan): implement runtime.Goexit correctly
	"text/scanner",
	"unicode",

	// Too slow:
	// "container/ring",
	// "hash/adler32",

	// TODO(adonovan): packages with Examples require os.Pipe (unimplemented):
	// "hash/crc32",
	// "unicode/utf8",
	// "log",
	// "path",
	// "flag",
}

type successPredicate func(exitcode int, output string) error

func run(t *testing.T, dir, input string, success successPredicate) bool {
	fmt.Printf("Input: %s\n", input)

	start := time.Now()

	var inputs []string
	for _, i := range strings.Split(input, " ") {
		if strings.HasSuffix(i, ".go") {
			i = dir + i
		}
		inputs = append(inputs, i)
	}

	var conf loader.Config
	if _, err := conf.FromArgs(inputs, true); err != nil {
		t.Errorf("FromArgs(%s) failed: %s", inputs, err)
		return false
	}

	conf.Import("runtime")

	// Print a helpful hint if we don't make it to the end.
	var hint string
	defer func() {
		if hint != "" {
			fmt.Println("FAIL")
			fmt.Println(hint)
		} else {
			fmt.Println("PASS")
		}

		interp.CapturedOutput = nil
	}()

	hint = fmt.Sprintf("To dump SSA representation, run:\n%% go build golang.org/x/tools/cmd/ssadump && ./ssadump -build=CFP %s\n", input)

	iprog, err := conf.Load()
	if err != nil {
		t.Errorf("conf.Load(%s) failed: %s", inputs, err)
		return false
	}

	prog := ssautil.CreateProgram(iprog, ssa.SanityCheckFunctions)
	prog.BuildAll()

	var mainPkg *ssa.Package
	var initialPkgs []*ssa.Package
	for _, info := range iprog.InitialPackages() {
		if info.Pkg.Path() == "runtime" {
			continue // not an initial package
		}
		p := prog.Package(info.Pkg)
		initialPkgs = append(initialPkgs, p)
		if mainPkg == nil && p.Func("main") != nil {
			mainPkg = p
		}
	}
	if mainPkg == nil {
		testmainPkg := prog.CreateTestMainPackage(initialPkgs...)
		if testmainPkg == nil {
			t.Errorf("CreateTestMainPackage(%s) returned nil", mainPkg)
			return false
		}
		if testmainPkg.Func("main") == nil {
			t.Errorf("synthetic testmain package has no main")
			return false
		}
		mainPkg = testmainPkg
	}

	var out bytes.Buffer
	interp.CapturedOutput = &out

	hint = fmt.Sprintf("To trace execution, run:\n%% go build golang.org/x/tools/cmd/ssadump && ./ssadump -build=C -run --interp=T %s\n", input)
	exitCode := interp.Interpret(mainPkg, 0, &types.StdSizes{8, 8}, inputs[0], []string{})

	// The definition of success varies with each file.
	if err := success(exitCode, out.String()); err != nil {
		t.Errorf("interp.Interpret(%s) failed: %s", inputs, err)
		return false
	}

	hint = "" // call off the hounds

	if false {
		fmt.Println(input, time.Since(start)) // test profiling
	}

	return true
}

const slash = string(os.PathSeparator)

func printFailures(failures []string) {
	if failures != nil {
		fmt.Println("The following tests failed:")
		for _, f := range failures {
			fmt.Printf("\t%s\n", f)
		}
	}
}

func success(exitcode int, output string) error {
	if exitcode != 0 {
		return fmt.Errorf("exit code was %d", exitcode)
	}
	if strings.Contains(output, "BUG") {
		return fmt.Errorf("exited zero but output contained 'BUG'")
	}
	return nil
}

// TestTestdataFiles runs the interpreter on testdata/*.go.
func TestTestdataFiles(t *testing.T) {
	var failures []string
	for _, input := range testdataTests {
		if !run(t, "testdata"+slash, input, success) {
			failures = append(failures, input)
		}
	}
	printFailures(failures)
}

// TestGorootTest runs the interpreter on $GOROOT/test/*.go.
func TestGorootTest(t *testing.T) {
	if testing.Short() {
		return // too slow (~30s)
	}

	var failures []string

	for _, input := range gorootTestTests {
		if !run(t, filepath.Join(build.Default.GOROOT, "test")+slash, input, success) {
			failures = append(failures, input)
		}
	}
	for _, input := range gorootSrcTests {
		if !run(t, filepath.Join(build.Default.GOROOT, "src")+slash, input, success) {
			failures = append(failures, input)
		}
	}
	printFailures(failures)
}

// TestTestmainPackage runs the interpreter on a synthetic "testmain" package.
func TestTestmainPackage(t *testing.T) {
	success := func(exitcode int, output string) error {
		if exitcode == 0 {
			return fmt.Errorf("unexpected success")
		}
		if !strings.Contains(output, "FAIL: TestFoo") {
			return fmt.Errorf("missing failure log for TestFoo")
		}
		if !strings.Contains(output, "FAIL: TestBar") {
			return fmt.Errorf("missing failure log for TestBar")
		}
		// TODO(adonovan): test benchmarks too
		return nil
	}
	run(t, "testdata"+slash, "a_test.go", success)
}

// CreateTestMainPackage should return nil if there were no tests.
func TestNullTestmainPackage(t *testing.T) {
	var conf loader.Config
	conf.CreateFromFilenames("", "testdata/b_test.go")
	iprog, err := conf.Load()
	if err != nil {
		t.Fatalf("CreatePackages failed: %s", err)
	}
	prog := ssautil.CreateProgram(iprog, ssa.SanityCheckFunctions)
	mainPkg := prog.Package(iprog.Created[0].Pkg)
	if mainPkg.Func("main") != nil {
		t.Fatalf("unexpected main function")
	}
	if prog.CreateTestMainPackage(mainPkg) != nil {
		t.Fatalf("CreateTestMainPackage returned non-nil")
	}
}
<|endoftext|>


<commit_before>// Copyright 2015, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package servenv

import (
	"crypto/tls"
	"crypto/x509"
	"flag"
	"fmt"
	"io/ioutil"
	"net"

	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"

	log "github.com/golang/glog"
)

// This file handles gRPC server, on its own port.
// Clients register servers, based on service map:
//
// servenv.RegisterGRPCFlags()
// servenv.OnRun(func() {
//   if servenv.GRPCCheckServiceMap("XXX") {
//     pb.RegisterXXX(servenv.GRPCServer, XXX)
//   }
// }
//
// Note servenv.GRPCServer can only be used in servenv.OnRun,
// and not before, as it is initialized right before calling OnRun.
var (
	// GRPCPort is the port to listen on for gRPC. If not set or zero, don't listen.
	GRPCPort *int

	// GRPCCert is the cert to use if TLS is enabled
	GRPCCert *string

	// GRPCKey is the key to use if TLS is enabled
	GRPCKey *string

	// GRPCCA is the CA to use if TLS is enabled
	GRPCCA *string

	// GRPCServer is the global server to serve gRPC.
	GRPCServer *grpc.Server
)

// isGRPCEnabled returns true if gRPC server is set
func isGRPCEnabled() bool {
	if GRPCPort != nil && *GRPCPort != 0 {
		return true
	}

	if SocketFile != nil && *SocketFile != "" {
		return true
	}

	return false
}

// createGRPCServer create the gRPC server we will be using.
// It has to be called after flags are parsed, but before
// services register themselves.
func createGRPCServer() {
	// skip if not registered
	if !isGRPCEnabled() {
		log.Infof("Skipping gRPC server creation")
		return
	}

	var opts []grpc.ServerOption
	if GRPCPort != nil && *GRPCCert != "" && *GRPCKey != "" {
		config := &tls.Config{}

		// load the server cert and key
		cert, err := tls.LoadX509KeyPair(*GRPCCert, *GRPCKey)
		if err != nil {
			log.Fatalf("Failed to load cert/key: %v", err)
		}
		config.Certificates = []tls.Certificate{cert}

		// if specified, load ca to validate client,
		// and enforce clients present valid certs.
		if *GRPCCA != "" {
			b, err := ioutil.ReadFile(*GRPCCA)
			if err != nil {
				log.Fatalf("Failed to read ca file: %v", err)
			}
			cp := x509.NewCertPool()
			if !cp.AppendCertsFromPEM(b) {
				log.Fatalf("Failed to append certificates")
			}
			config.ClientCAs = cp
			config.ClientAuth = tls.RequireAndVerifyClientCert
		}

		// create the creds server options
		creds := credentials.NewTLS(config)
		opts = []grpc.ServerOption{grpc.Creds(creds)}
	}

	GRPCServer = grpc.NewServer(opts...)
}

func serveGRPC() {
	// skip if not registered
	if GRPCPort == nil || *GRPCPort == 0 {
		return
	}

	// listen on the port
	log.Infof("Listening for gRPC calls on port %v", *GRPCPort)
	listener, err := net.Listen("tcp", fmt.Sprintf(":%d", *GRPCPort))
	if err != nil {
		log.Fatalf("Cannot listen on port %v for gRPC: %v", *GRPCPort, err)
	}

	// and serve on it
	go GRPCServer.Serve(listener)
}

// RegisterGRPCFlags registers the right command line flag to enable gRPC
func RegisterGRPCFlags() {
	GRPCPort = flag.Int("grpc_port", 0, "Port to listen on for gRPC calls")
	GRPCCert = flag.String("grpc_cert", "", "certificate to use, requires grpc_key, enables TLS")
	GRPCKey = flag.String("grpc_key", "", "key to use, requires grpc_cert, enables TLS")
	GRPCCA = flag.String("grpc_ca", "", "ca to use, requires TLS, and enforces client cert check")
}

// GRPCCheckServiceMap returns if we should register a gRPC service
// (and also logs how to enable / disable it)
func GRPCCheckServiceMap(name string) bool {
	// Silently fail individual services if gRPC is not enabled in
	// the first place (either on a grpc port or on the socket file)
	if !isGRPCEnabled() {
		return false
	}

	// then check ServiceMap
	return CheckServiceMap("grpc", name)
}
<commit_msg>Allow to override the default gRPC max message size.<commit_after>// Copyright 2015, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package servenv

import (
	"crypto/tls"
	"crypto/x509"
	"flag"
	"fmt"
	"io/ioutil"
	"net"

	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"

	log "github.com/golang/glog"
)

// This file handles gRPC server, on its own port.
// Clients register servers, based on service map:
//
// servenv.RegisterGRPCFlags()
// servenv.OnRun(func() {
//   if servenv.GRPCCheckServiceMap("XXX") {
//     pb.RegisterXXX(servenv.GRPCServer, XXX)
//   }
// }
//
// Note servenv.GRPCServer can only be used in servenv.OnRun,
// and not before, as it is initialized right before calling OnRun.
var (
	// GRPCPort is the port to listen on for gRPC. If not set or zero, don't listen.
	GRPCPort *int

	// GRPCCert is the cert to use if TLS is enabled
	GRPCCert *string

	// GRPCKey is the key to use if TLS is enabled
	GRPCKey *string

	// GRPCCA is the CA to use if TLS is enabled
	GRPCCA *string

	// GRPCMaxMessageSize is the maximum message size which the gRPC server will
	// accept. Larger messages will be rejected.
	GRPCMaxMessageSize *int

	// GRPCServer is the global server to serve gRPC.
	GRPCServer *grpc.Server
)

// isGRPCEnabled returns true if gRPC server is set
func isGRPCEnabled() bool {
	if GRPCPort != nil && *GRPCPort != 0 {
		return true
	}

	if SocketFile != nil && *SocketFile != "" {
		return true
	}

	return false
}

// createGRPCServer create the gRPC server we will be using.
// It has to be called after flags are parsed, but before
// services register themselves.
func createGRPCServer() {
	// skip if not registered
	if !isGRPCEnabled() {
		log.Infof("Skipping gRPC server creation")
		return
	}

	var opts []grpc.ServerOption
	if GRPCPort != nil && *GRPCCert != "" && *GRPCKey != "" {
		config := &tls.Config{}

		// load the server cert and key
		cert, err := tls.LoadX509KeyPair(*GRPCCert, *GRPCKey)
		if err != nil {
			log.Fatalf("Failed to load cert/key: %v", err)
		}
		config.Certificates = []tls.Certificate{cert}

		// if specified, load ca to validate client,
		// and enforce clients present valid certs.
		if *GRPCCA != "" {
			b, err := ioutil.ReadFile(*GRPCCA)
			if err != nil {
				log.Fatalf("Failed to read ca file: %v", err)
			}
			cp := x509.NewCertPool()
			if !cp.AppendCertsFromPEM(b) {
				log.Fatalf("Failed to append certificates")
			}
			config.ClientCAs = cp
			config.ClientAuth = tls.RequireAndVerifyClientCert
		}

		// create the creds server options
		creds := credentials.NewTLS(config)
		opts = []grpc.ServerOption{grpc.Creds(creds)}
	}
	// Override the default max message size (which is 4 MiB in gRPC 1.0.0).
	// Large messages can occur when users try to insert very big rows. If they
	// hit the limit, they'll see the following error:
	// grpc: received message length XXXXXXX exceeding the max size 4194304
	// Note: For gRPC 1.0.0 it's sufficient to set the limit on the server only
	// because it's not enforced on the client side.
	if GRPCMaxMessageSize != nil {
		opts = append(opts, grpc.MaxMsgSize(*GRPCMaxMessageSize))
	}

	GRPCServer = grpc.NewServer(opts...)
}

func serveGRPC() {
	// skip if not registered
	if GRPCPort == nil || *GRPCPort == 0 {
		return
	}

	// listen on the port
	log.Infof("Listening for gRPC calls on port %v", *GRPCPort)
	listener, err := net.Listen("tcp", fmt.Sprintf(":%d", *GRPCPort))
	if err != nil {
		log.Fatalf("Cannot listen on port %v for gRPC: %v", *GRPCPort, err)
	}

	// and serve on it
	go GRPCServer.Serve(listener)
}

// RegisterGRPCFlags registers the right command line flag to enable gRPC
func RegisterGRPCFlags() {
	GRPCPort = flag.Int("grpc_port", 0, "Port to listen on for gRPC calls")
	GRPCCert = flag.String("grpc_cert", "", "certificate to use, requires grpc_key, enables TLS")
	GRPCKey = flag.String("grpc_key", "", "key to use, requires grpc_cert, enables TLS")
	GRPCCA = flag.String("grpc_ca", "", "ca to use, requires TLS, and enforces client cert check")
	// Note: We're using 4 MiB as default value because that's the default in the
	// gRPC 1.0.0 Go server.
	GRPCMaxMessageSize = flag.Int("grpc_max_message_size", 4*1024*1024, "Maximum allowed RPC message size. Larger messages will be rejected by gRPC with the error 'exceeding the max size'.")
}

// GRPCCheckServiceMap returns if we should register a gRPC service
// (and also logs how to enable / disable it)
func GRPCCheckServiceMap(name string) bool {
	// Silently fail individual services if gRPC is not enabled in
	// the first place (either on a grpc port or on the socket file)
	if !isGRPCEnabled() {
		return false
	}

	// then check ServiceMap
	return CheckServiceMap("grpc", name)
}
<|endoftext|>


<commit_before>// Copyright 2012, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

/*
The agent handles local execution of actions triggered remotely.
It has two execution models:

- listening on an action path for ActionNode objects.  When receiving
  an action, it will forward it to vtaction to perform it (vtaction
  uses the actor code). We usually use this model for long-running
  queries where an RPC would time out.

  All vtaction calls lock the actionMutex.

  After executing vtaction, we always call agent.changeCallback.
  Additionnally, for TABLET_ACTION_APPLY_SCHEMA, we will force a schema
  reload.

- listening as an RPC server. The agent performs the action itself,
  calling the actor code directly. We use this for short lived actions.

  Most RPC calls lock the actionMutex, except the easy read-donly ones.

  We will not call changeCallback for all actions, just for the ones
  that are relevant. Same for schema reload.

  See rpc_server.go for all cases, and which action takes the actionMutex,
  runs changeCallback, and reloads the schema.

*/

package tabletmanager

import (
	"encoding/json"
	"flag"
	"fmt"
	"net"
	"os"
	"os/exec"
	"path"
	"sync"
	"time"

	log "github.com/golang/glog"
	"github.com/youtube/vitess/go/history"
	"github.com/youtube/vitess/go/jscfg"
	"github.com/youtube/vitess/go/netutil"
	"github.com/youtube/vitess/go/stats"
	"github.com/youtube/vitess/go/vt/dbconfigs"
	"github.com/youtube/vitess/go/vt/env"
	"github.com/youtube/vitess/go/vt/logutil"
	"github.com/youtube/vitess/go/vt/mysqlctl"
	"github.com/youtube/vitess/go/vt/tabletmanager/actionnode"
	"github.com/youtube/vitess/go/vt/tabletmanager/actor"
	"github.com/youtube/vitess/go/vt/tabletserver"
	"github.com/youtube/vitess/go/vt/topo"
)

var (
	vtactionBinaryPath = flag.String("vtaction_binary_path", "", "Full path (including filename) to vtaction binary. If not set, tries VTROOT/bin/vtaction.")
)

type tabletChangeItem struct {
	oldTablet  topo.Tablet
	newTablet  topo.Tablet
	context    string
	queuedTime time.Time
}

// ActionAgent is the main class for the agent.
type ActionAgent struct {
	// The following fields are set during creation
	TopoServer      topo.Server
	TabletAlias     topo.TabletAlias
	Mysqld          *mysqlctl.Mysqld
	DBConfigs       *dbconfigs.DBConfigs
	SchemaOverrides []tabletserver.SchemaOverride
	BinlogPlayerMap *BinlogPlayerMap

	// Internal variables
	vtActionBinFile string        // path to vtaction binary
	done            chan struct{} // closed when we are done.

	// This is the History of the health checks, public so status
	// pages can display it
	History            *history.History
	lastHealthMapCount *stats.Int

	// actionMutex is there to run only one action at a time. If
	// both agent.actionMutex and agent.mutex needs to be taken,
	// take actionMutex first.
	actionMutex sync.Mutex // to run only one action at a time

	// mutex protects _tablet and serializes writes to changeItems.
	mutex       sync.Mutex
	changeItems chan tabletChangeItem
	_tablet     *topo.TabletInfo
}

func loadSchemaOverrides(overridesFile string) []tabletserver.SchemaOverride {
	var schemaOverrides []tabletserver.SchemaOverride
	if overridesFile == "" {
		return schemaOverrides
	}
	if err := jscfg.ReadJson(overridesFile, &schemaOverrides); err != nil {
		log.Warningf("can't read overridesFile %v: %v", overridesFile, err)
	} else {
		data, _ := json.MarshalIndent(schemaOverrides, "", "  ")
		log.Infof("schemaOverrides: %s\n", data)
	}
	return schemaOverrides
}

// NewActionAgent creates a new ActionAgent and registers all the
// associated services
func NewActionAgent(
	tabletAlias topo.TabletAlias,
	dbcfgs *dbconfigs.DBConfigs,
	mycnf *mysqlctl.Mycnf,
	port, securePort int,
	overridesFile string,
) (agent *ActionAgent, err error) {
	schemaOverrides := loadSchemaOverrides(overridesFile)

	topoServer := topo.GetServer()
	mysqld := mysqlctl.NewMysqld("Dba", mycnf, &dbcfgs.Dba, &dbcfgs.Repl)

	agent = &ActionAgent{
		TopoServer:         topoServer,
		TabletAlias:        tabletAlias,
		Mysqld:             mysqld,
		DBConfigs:          dbcfgs,
		SchemaOverrides:    schemaOverrides,
		done:               make(chan struct{}),
		History:            history.New(historyLength),
		lastHealthMapCount: stats.NewInt("LastHealthMapCount"),
		changeItems:        make(chan tabletChangeItem, 100),
	}

	// Start the binlog player services, not playing at start.
	agent.BinlogPlayerMap = NewBinlogPlayerMap(topoServer, &dbcfgs.App.ConnectionParams, mysqld)
	RegisterBinlogPlayerMap(agent.BinlogPlayerMap)

	// try to figure out the mysql port
	mysqlPort := mycnf.MysqlPort
	if mysqlPort == 0 {
		// we don't know the port, try to get it from mysqld
		var err error
		mysqlPort, err = mysqld.GetMysqlPort()
		if err != nil {
			log.Warningf("Cannot get current mysql port, will use 0 for now: %v", err)
		}
	}

	if err := agent.Start(mysqlPort, port, securePort); err != nil {
		return nil, err
	}

	// register the RPC services from the agent
	agent.registerQueryService()

	// start health check if needed
	agent.initHeathCheck()

	return agent, nil
}

func (agent *ActionAgent) runChangeCallback(oldTablet *topo.Tablet, context string) {
	agent.mutex.Lock()
	// Access directly since we have the lock.
	newTablet := agent._tablet.Tablet
	agent.changeItems <- tabletChangeItem{oldTablet: *oldTablet, newTablet: *newTablet, context: context, queuedTime: time.Now()}
	log.Infof("Queued tablet callback: %v", context)
	agent.mutex.Unlock()
}

func (agent *ActionAgent) executeCallbacksLoop() {
	for {
		select {
		case changeItem := <-agent.changeItems:
			log.Infof("Running tablet callback after %v: %v", time.Now().Sub(changeItem.queuedTime), changeItem.context)
			agent.changeCallback(changeItem.oldTablet, changeItem.newTablet)
		case <-agent.done:
			return
		}
	}
}

func (agent *ActionAgent) readTablet() error {
	tablet, err := agent.TopoServer.GetTablet(agent.TabletAlias)
	if err != nil {
		return err
	}
	agent.mutex.Lock()
	agent._tablet = tablet
	agent.mutex.Unlock()
	return nil
}

func (agent *ActionAgent) Tablet() *topo.TabletInfo {
	agent.mutex.Lock()
	tablet := agent._tablet
	agent.mutex.Unlock()
	return tablet
}

func (agent *ActionAgent) resolvePaths() error {
	var p string
	if *vtactionBinaryPath != "" {
		p = *vtactionBinaryPath
	} else {
		vtroot, err := env.VtRoot()
		if err != nil {
			return err
		}
		p = path.Join(vtroot, "bin/vtaction")
	}
	if _, err := os.Stat(p); err != nil {
		return fmt.Errorf("vtaction binary %s not found: %v", p, err)
	}
	agent.vtActionBinFile = p
	return nil
}

// A non-nil return signals that event processing should stop.
func (agent *ActionAgent) dispatchAction(actionPath, data string) error {
	agent.actionMutex.Lock()
	defer agent.actionMutex.Unlock()

	log.Infof("action dispatch %v", actionPath)
	actionNode, err := actionnode.ActionNodeFromJson(data, actionPath)
	if err != nil {
		log.Errorf("action decode failed: %v %v", actionPath, err)
		return nil
	}

	cmd := []string{
		agent.vtActionBinFile,
		"-action", actionNode.Action,
		"-action-node", actionPath,
		"-action-guid", actionNode.ActionGuid,
	}
	cmd = append(cmd, logutil.GetSubprocessFlags()...)
	cmd = append(cmd, topo.GetSubprocessFlags()...)
	cmd = append(cmd, dbconfigs.GetSubprocessFlags()...)
	cmd = append(cmd, mysqlctl.GetSubprocessFlags()...)
	log.Infof("action launch %v", cmd)
	vtActionCmd := exec.Command(cmd[0], cmd[1:]...)

	stdOut, vtActionErr := vtActionCmd.CombinedOutput()
	if vtActionErr != nil {
		log.Errorf("agent action failed: %v %v\n%s", actionPath, vtActionErr, stdOut)
		// If the action failed, preserve single execution path semantics.
		return vtActionErr
	}

	log.Infof("Agent action completed %v %s", actionPath, stdOut)
	agent.afterAction(actionPath, actionNode.Action == actionnode.TABLET_ACTION_APPLY_SCHEMA)
	return nil
}

// afterAction needs to be run after an action may have changed the current
// state of the tablet.
func (agent *ActionAgent) afterAction(context string, reloadSchema bool) {
	log.Infof("Executing post-action change callbacks")

	// Save the old tablet so callbacks can have a better idea of
	// the precise nature of the transition.
	oldTablet := agent.Tablet().Tablet

	// Actions should have side effects on the tablet, so reload the data.
	if err := agent.readTablet(); err != nil {
		log.Warningf("Failed rereading tablet after %v - services may be inconsistent: %v", context, err)
	} else {
		if updatedTablet := actor.CheckTabletMysqlPort(agent.TopoServer, agent.Mysqld, agent.Tablet()); updatedTablet != nil {
			agent.mutex.Lock()
			agent._tablet = updatedTablet
			agent.mutex.Unlock()
		}

		agent.runChangeCallback(oldTablet, context)
	}

	// Maybe invalidate the schema.
	// This adds a dependency between tabletmanager and tabletserver,
	// so it's not ideal. But I (alainjobart) think it's better
	// to have up to date schema in vtocc.
	if reloadSchema {
		tabletserver.ReloadSchema()
	}
	log.Infof("Done with post-action change callbacks")
}

func (agent *ActionAgent) verifyTopology() error {
	tablet := agent.Tablet()
	if tablet == nil {
		return fmt.Errorf("agent._tablet is nil")
	}

	if err := topo.Validate(agent.TopoServer, agent.TabletAlias); err != nil {
		// Don't stop, it's not serious enough, this is likely transient.
		log.Warningf("tablet validate failed: %v %v", agent.TabletAlias, err)
	}

	return agent.TopoServer.ValidateTabletActions(agent.TabletAlias)
}

func (agent *ActionAgent) verifyServingAddrs() error {
	if !agent.Tablet().IsRunningQueryService() {
		return nil
	}

	// Check to see our address is registered in the right place.
	addr, err := agent.Tablet().Tablet.EndPoint()
	if err != nil {
		return err
	}
	return agent.TopoServer.UpdateTabletEndpoint(agent.Tablet().Tablet.Alias.Cell, agent.Tablet().Keyspace, agent.Tablet().Shard, agent.Tablet().Type, addr)
}

// bindAddr: the address for the query service advertised by this agent
func (agent *ActionAgent) Start(mysqlPort, vtPort, vtsPort int) error {
	var err error
	if err = agent.readTablet(); err != nil {
		return err
	}

	if err = agent.resolvePaths(); err != nil {
		return err
	}

	// find our hostname as fully qualified, and IP
	hostname, err := netutil.FullyQualifiedHostname()
	if err != nil {
		return err
	}
	ipAddrs, err := net.LookupHost(hostname)
	if err != nil {
		return err
	}
	ipAddr := ipAddrs[0]

	// Update bind addr for mysql and query service in the tablet node.
	f := func(tablet *topo.Tablet) error {
		tablet.Hostname = hostname
		tablet.IPAddr = ipAddr
		if tablet.Portmap == nil {
			tablet.Portmap = make(map[string]int)
		}
		if mysqlPort != 0 {
			// only overwrite mysql port if we know it, otherwise
			// leave it as is.
			tablet.Portmap["mysql"] = mysqlPort
		}
		tablet.Portmap["vt"] = vtPort
		if vtsPort != 0 {
			tablet.Portmap["vts"] = vtsPort
		} else {
			delete(tablet.Portmap, "vts")
		}
		return nil
	}
	if err := agent.TopoServer.UpdateTabletFields(agent.Tablet().Alias, f); err != nil {
		return err
	}

	// Reread to get the changes we just made
	if err := agent.readTablet(); err != nil {
		return err
	}

	data := fmt.Sprintf("host:%v\npid:%v\n", hostname, os.Getpid())

	if err := agent.TopoServer.CreateTabletPidNode(agent.TabletAlias, data, agent.done); err != nil {
		return err
	}

	if err = agent.verifyTopology(); err != nil {
		return err
	}

	if err = agent.verifyServingAddrs(); err != nil {
		return err
	}

	oldTablet := &topo.Tablet{}
	agent.runChangeCallback(oldTablet, "Start")

	go agent.actionEventLoop()
	go agent.executeCallbacksLoop()
	return nil
}

func (agent *ActionAgent) Stop() {
	close(agent.done)
	agent.BinlogPlayerMap.StopAllPlayersAndReset()
	agent.Mysqld.Close()
}

func (agent *ActionAgent) actionEventLoop() {
	f := func(actionPath, data string) error {
		return agent.dispatchAction(actionPath, data)
	}
	agent.TopoServer.ActionEventLoop(agent.TabletAlias, f, agent.done)
}
<commit_msg>Allow plugin files to add subprocess flags for site-local packages.<commit_after>// Copyright 2012, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

/*
The agent handles local execution of actions triggered remotely.
It has two execution models:

- listening on an action path for ActionNode objects.  When receiving
  an action, it will forward it to vtaction to perform it (vtaction
  uses the actor code). We usually use this model for long-running
  queries where an RPC would time out.

  All vtaction calls lock the actionMutex.

  After executing vtaction, we always call agent.changeCallback.
  Additionnally, for TABLET_ACTION_APPLY_SCHEMA, we will force a schema
  reload.

- listening as an RPC server. The agent performs the action itself,
  calling the actor code directly. We use this for short lived actions.

  Most RPC calls lock the actionMutex, except the easy read-donly ones.

  We will not call changeCallback for all actions, just for the ones
  that are relevant. Same for schema reload.

  See rpc_server.go for all cases, and which action takes the actionMutex,
  runs changeCallback, and reloads the schema.

*/

package tabletmanager

import (
	"encoding/json"
	"flag"
	"fmt"
	"net"
	"os"
	"os/exec"
	"path"
	"sync"
	"time"

	log "github.com/golang/glog"
	"github.com/youtube/vitess/go/history"
	"github.com/youtube/vitess/go/jscfg"
	"github.com/youtube/vitess/go/netutil"
	"github.com/youtube/vitess/go/stats"
	"github.com/youtube/vitess/go/vt/dbconfigs"
	"github.com/youtube/vitess/go/vt/env"
	"github.com/youtube/vitess/go/vt/logutil"
	"github.com/youtube/vitess/go/vt/mysqlctl"
	"github.com/youtube/vitess/go/vt/tabletmanager/actionnode"
	"github.com/youtube/vitess/go/vt/tabletmanager/actor"
	"github.com/youtube/vitess/go/vt/tabletserver"
	"github.com/youtube/vitess/go/vt/topo"
)

var (
	vtactionBinaryPath = flag.String("vtaction_binary_path", "", "Full path (including filename) to vtaction binary. If not set, tries VTROOT/bin/vtaction.")
)

type tabletChangeItem struct {
	oldTablet  topo.Tablet
	newTablet  topo.Tablet
	context    string
	queuedTime time.Time
}

// ActionAgent is the main class for the agent.
type ActionAgent struct {
	// The following fields are set during creation
	TopoServer      topo.Server
	TabletAlias     topo.TabletAlias
	Mysqld          *mysqlctl.Mysqld
	DBConfigs       *dbconfigs.DBConfigs
	SchemaOverrides []tabletserver.SchemaOverride
	BinlogPlayerMap *BinlogPlayerMap

	// Internal variables
	vtActionBinFile string        // path to vtaction binary
	done            chan struct{} // closed when we are done.

	// This is the History of the health checks, public so status
	// pages can display it
	History            *history.History
	lastHealthMapCount *stats.Int

	// actionMutex is there to run only one action at a time. If
	// both agent.actionMutex and agent.mutex needs to be taken,
	// take actionMutex first.
	actionMutex sync.Mutex // to run only one action at a time

	// mutex protects _tablet and serializes writes to changeItems.
	mutex       sync.Mutex
	changeItems chan tabletChangeItem
	_tablet     *topo.TabletInfo
}

func loadSchemaOverrides(overridesFile string) []tabletserver.SchemaOverride {
	var schemaOverrides []tabletserver.SchemaOverride
	if overridesFile == "" {
		return schemaOverrides
	}
	if err := jscfg.ReadJson(overridesFile, &schemaOverrides); err != nil {
		log.Warningf("can't read overridesFile %v: %v", overridesFile, err)
	} else {
		data, _ := json.MarshalIndent(schemaOverrides, "", "  ")
		log.Infof("schemaOverrides: %s\n", data)
	}
	return schemaOverrides
}

// NewActionAgent creates a new ActionAgent and registers all the
// associated services
func NewActionAgent(
	tabletAlias topo.TabletAlias,
	dbcfgs *dbconfigs.DBConfigs,
	mycnf *mysqlctl.Mycnf,
	port, securePort int,
	overridesFile string,
) (agent *ActionAgent, err error) {
	schemaOverrides := loadSchemaOverrides(overridesFile)

	topoServer := topo.GetServer()
	mysqld := mysqlctl.NewMysqld("Dba", mycnf, &dbcfgs.Dba, &dbcfgs.Repl)

	agent = &ActionAgent{
		TopoServer:         topoServer,
		TabletAlias:        tabletAlias,
		Mysqld:             mysqld,
		DBConfigs:          dbcfgs,
		SchemaOverrides:    schemaOverrides,
		done:               make(chan struct{}),
		History:            history.New(historyLength),
		lastHealthMapCount: stats.NewInt("LastHealthMapCount"),
		changeItems:        make(chan tabletChangeItem, 100),
	}

	// Start the binlog player services, not playing at start.
	agent.BinlogPlayerMap = NewBinlogPlayerMap(topoServer, &dbcfgs.App.ConnectionParams, mysqld)
	RegisterBinlogPlayerMap(agent.BinlogPlayerMap)

	// try to figure out the mysql port
	mysqlPort := mycnf.MysqlPort
	if mysqlPort == 0 {
		// we don't know the port, try to get it from mysqld
		var err error
		mysqlPort, err = mysqld.GetMysqlPort()
		if err != nil {
			log.Warningf("Cannot get current mysql port, will use 0 for now: %v", err)
		}
	}

	if err := agent.Start(mysqlPort, port, securePort); err != nil {
		return nil, err
	}

	// register the RPC services from the agent
	agent.registerQueryService()

	// start health check if needed
	agent.initHeathCheck()

	return agent, nil
}

func (agent *ActionAgent) runChangeCallback(oldTablet *topo.Tablet, context string) {
	agent.mutex.Lock()
	// Access directly since we have the lock.
	newTablet := agent._tablet.Tablet
	agent.changeItems <- tabletChangeItem{oldTablet: *oldTablet, newTablet: *newTablet, context: context, queuedTime: time.Now()}
	log.Infof("Queued tablet callback: %v", context)
	agent.mutex.Unlock()
}

func (agent *ActionAgent) executeCallbacksLoop() {
	for {
		select {
		case changeItem := <-agent.changeItems:
			log.Infof("Running tablet callback after %v: %v", time.Now().Sub(changeItem.queuedTime), changeItem.context)
			agent.changeCallback(changeItem.oldTablet, changeItem.newTablet)
		case <-agent.done:
			return
		}
	}
}

func (agent *ActionAgent) readTablet() error {
	tablet, err := agent.TopoServer.GetTablet(agent.TabletAlias)
	if err != nil {
		return err
	}
	agent.mutex.Lock()
	agent._tablet = tablet
	agent.mutex.Unlock()
	return nil
}

func (agent *ActionAgent) Tablet() *topo.TabletInfo {
	agent.mutex.Lock()
	tablet := agent._tablet
	agent.mutex.Unlock()
	return tablet
}

func (agent *ActionAgent) resolvePaths() error {
	var p string
	if *vtactionBinaryPath != "" {
		p = *vtactionBinaryPath
	} else {
		vtroot, err := env.VtRoot()
		if err != nil {
			return err
		}
		p = path.Join(vtroot, "bin/vtaction")
	}
	if _, err := os.Stat(p); err != nil {
		return fmt.Errorf("vtaction binary %s not found: %v", p, err)
	}
	agent.vtActionBinFile = p
	return nil
}

// A non-nil return signals that event processing should stop.
func (agent *ActionAgent) dispatchAction(actionPath, data string) error {
	agent.actionMutex.Lock()
	defer agent.actionMutex.Unlock()

	log.Infof("action dispatch %v", actionPath)
	actionNode, err := actionnode.ActionNodeFromJson(data, actionPath)
	if err != nil {
		log.Errorf("action decode failed: %v %v", actionPath, err)
		return nil
	}

	cmd := []string{
		agent.vtActionBinFile,
		"-action", actionNode.Action,
		"-action-node", actionPath,
		"-action-guid", actionNode.ActionGuid,
	}
	for _, getSubprocessFlags := range getSubprocessFlagsFuncs {
		cmd = append(cmd, getSubprocessFlags()...)
	}
	log.Infof("action launch %v", cmd)
	vtActionCmd := exec.Command(cmd[0], cmd[1:]...)

	stdOut, vtActionErr := vtActionCmd.CombinedOutput()
	if vtActionErr != nil {
		log.Errorf("agent action failed: %v %v\n%s", actionPath, vtActionErr, stdOut)
		// If the action failed, preserve single execution path semantics.
		return vtActionErr
	}

	log.Infof("Agent action completed %v %s", actionPath, stdOut)
	agent.afterAction(actionPath, actionNode.Action == actionnode.TABLET_ACTION_APPLY_SCHEMA)
	return nil
}

// afterAction needs to be run after an action may have changed the current
// state of the tablet.
func (agent *ActionAgent) afterAction(context string, reloadSchema bool) {
	log.Infof("Executing post-action change callbacks")

	// Save the old tablet so callbacks can have a better idea of
	// the precise nature of the transition.
	oldTablet := agent.Tablet().Tablet

	// Actions should have side effects on the tablet, so reload the data.
	if err := agent.readTablet(); err != nil {
		log.Warningf("Failed rereading tablet after %v - services may be inconsistent: %v", context, err)
	} else {
		if updatedTablet := actor.CheckTabletMysqlPort(agent.TopoServer, agent.Mysqld, agent.Tablet()); updatedTablet != nil {
			agent.mutex.Lock()
			agent._tablet = updatedTablet
			agent.mutex.Unlock()
		}

		agent.runChangeCallback(oldTablet, context)
	}

	// Maybe invalidate the schema.
	// This adds a dependency between tabletmanager and tabletserver,
	// so it's not ideal. But I (alainjobart) think it's better
	// to have up to date schema in vtocc.
	if reloadSchema {
		tabletserver.ReloadSchema()
	}
	log.Infof("Done with post-action change callbacks")
}

func (agent *ActionAgent) verifyTopology() error {
	tablet := agent.Tablet()
	if tablet == nil {
		return fmt.Errorf("agent._tablet is nil")
	}

	if err := topo.Validate(agent.TopoServer, agent.TabletAlias); err != nil {
		// Don't stop, it's not serious enough, this is likely transient.
		log.Warningf("tablet validate failed: %v %v", agent.TabletAlias, err)
	}

	return agent.TopoServer.ValidateTabletActions(agent.TabletAlias)
}

func (agent *ActionAgent) verifyServingAddrs() error {
	if !agent.Tablet().IsRunningQueryService() {
		return nil
	}

	// Check to see our address is registered in the right place.
	addr, err := agent.Tablet().Tablet.EndPoint()
	if err != nil {
		return err
	}
	return agent.TopoServer.UpdateTabletEndpoint(agent.Tablet().Tablet.Alias.Cell, agent.Tablet().Keyspace, agent.Tablet().Shard, agent.Tablet().Type, addr)
}

// bindAddr: the address for the query service advertised by this agent
func (agent *ActionAgent) Start(mysqlPort, vtPort, vtsPort int) error {
	var err error
	if err = agent.readTablet(); err != nil {
		return err
	}

	if err = agent.resolvePaths(); err != nil {
		return err
	}

	// find our hostname as fully qualified, and IP
	hostname, err := netutil.FullyQualifiedHostname()
	if err != nil {
		return err
	}
	ipAddrs, err := net.LookupHost(hostname)
	if err != nil {
		return err
	}
	ipAddr := ipAddrs[0]

	// Update bind addr for mysql and query service in the tablet node.
	f := func(tablet *topo.Tablet) error {
		tablet.Hostname = hostname
		tablet.IPAddr = ipAddr
		if tablet.Portmap == nil {
			tablet.Portmap = make(map[string]int)
		}
		if mysqlPort != 0 {
			// only overwrite mysql port if we know it, otherwise
			// leave it as is.
			tablet.Portmap["mysql"] = mysqlPort
		}
		tablet.Portmap["vt"] = vtPort
		if vtsPort != 0 {
			tablet.Portmap["vts"] = vtsPort
		} else {
			delete(tablet.Portmap, "vts")
		}
		return nil
	}
	if err := agent.TopoServer.UpdateTabletFields(agent.Tablet().Alias, f); err != nil {
		return err
	}

	// Reread to get the changes we just made
	if err := agent.readTablet(); err != nil {
		return err
	}

	data := fmt.Sprintf("host:%v\npid:%v\n", hostname, os.Getpid())

	if err := agent.TopoServer.CreateTabletPidNode(agent.TabletAlias, data, agent.done); err != nil {
		return err
	}

	if err = agent.verifyTopology(); err != nil {
		return err
	}

	if err = agent.verifyServingAddrs(); err != nil {
		return err
	}

	oldTablet := &topo.Tablet{}
	agent.runChangeCallback(oldTablet, "Start")

	go agent.actionEventLoop()
	go agent.executeCallbacksLoop()
	return nil
}

func (agent *ActionAgent) Stop() {
	close(agent.done)
	agent.BinlogPlayerMap.StopAllPlayersAndReset()
	agent.Mysqld.Close()
}

func (agent *ActionAgent) actionEventLoop() {
	f := func(actionPath, data string) error {
		return agent.dispatchAction(actionPath, data)
	}
	agent.TopoServer.ActionEventLoop(agent.TabletAlias, f, agent.done)
}

var getSubprocessFlagsFuncs []func() []string

func init() {
	getSubprocessFlagsFuncs = append(getSubprocessFlagsFuncs, logutil.GetSubprocessFlags)
	getSubprocessFlagsFuncs = append(getSubprocessFlagsFuncs, topo.GetSubprocessFlags)
	getSubprocessFlagsFuncs = append(getSubprocessFlagsFuncs, dbconfigs.GetSubprocessFlags)
	getSubprocessFlagsFuncs = append(getSubprocessFlagsFuncs, mysqlctl.GetSubprocessFlags)
}
<|endoftext|>


<commit_before>// Copyright 2015, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package tabletserver

import (
	"fmt"
	"time"

	log "github.com/golang/glog"
	mproto "github.com/youtube/vitess/go/mysql/proto"
	"github.com/youtube/vitess/go/sqldb"
	"github.com/youtube/vitess/go/sync2"
	"github.com/youtube/vitess/go/vt/dbconnpool"
	"golang.org/x/net/context"
)

// DBConn is a db connection for tabletserver.
// It performs automatic reconnects as needed.
// Its Execute function has a timeout that can kill
// its own queries and the underlying connection.
// It will also trigger a CheckMySQL whenever applicable.
type DBConn struct {
	conn *dbconnpool.DBConnection
	info *sqldb.ConnParams
	pool *ConnPool

	current sync2.AtomicString
}

// NewDBConn creates a new DBConn. It triggers a CheckMySQL if creation fails.
func NewDBConn(cp *ConnPool, appParams, dbaParams *sqldb.ConnParams) (*DBConn, error) {
	c, err := dbconnpool.NewDBConnection(appParams, mysqlStats)
	if err != nil {
		go checkMySQL()
		return nil, err
	}
	return &DBConn{
		conn: c,
		info: appParams,
		pool: cp,
	}, nil
}

// Exec executes the specified query. If there is a connection error, it will reconnect
// and retry. A failed reconnect will trigger a CheckMySQL.
func (dbc *DBConn) Exec(ctx context.Context, query string, maxrows int, wantfields bool) (*mproto.QueryResult, error) {
	for attempt := 1; attempt <= 2; attempt++ {
		r, err := dbc.execOnce(ctx, query, maxrows, wantfields)
		switch {
		case err == nil:
			return r, nil
		case !IsConnErr(err):
			return nil, NewTabletErrorSql(ErrFail, err)
		case attempt == 2:
			return nil, NewTabletErrorSql(ErrFatal, err)
		}
		err2 := dbc.reconnect()
		if err2 != nil {
			go checkMySQL()
			return nil, NewTabletErrorSql(ErrFatal, err)
		}
	}
	panic("unreachable")
}

func (dbc *DBConn) execOnce(ctx context.Context, query string, maxrows int, wantfields bool) (*mproto.QueryResult, error) {
	dbc.current.Set(query)
	defer dbc.current.Set("")

	done, err := dbc.setDeadline(ctx)
	if err != nil {
		return nil, err
	}
	if done != nil {
		defer close(done)
	}
	// Uncomment this line for manual testing.
	// defer time.Sleep(20 * time.Second)
	return dbc.conn.ExecuteFetch(query, maxrows, wantfields)
}

// ExecOnce executes the specified query, but does not retry on connection errors.
func (dbc *DBConn) ExecOnce(ctx context.Context, query string, maxrows int, wantfields bool) (*mproto.QueryResult, error) {
	return dbc.execOnce(ctx, query, maxrows, wantfields)
}

// Stream executes the query and streams the results.
func (dbc *DBConn) Stream(ctx context.Context, query string, callback func(*mproto.QueryResult) error, streamBufferSize int) error {
	dbc.current.Set(query)
	defer dbc.current.Set("")

	done, err := dbc.setDeadline(ctx)
	if err != nil {
		return err
	}
	if done != nil {
		defer close(done)
	}
	return dbc.conn.ExecuteStreamFetch(query, callback, streamBufferSize)
}

// VerifyStrict returns true if MySQL is in STRICT mode.
func (dbc *DBConn) VerifyStrict() bool {
	return dbc.conn.VerifyStrict()
}

// Close closes the DBConn.
func (dbc *DBConn) Close() {
	dbc.conn.Close()
}

// IsClosed returns true if DBConn is closed.
func (dbc *DBConn) IsClosed() bool {
	return dbc.conn.IsClosed()
}

// Recycle returns the DBConn to the pool.
func (dbc *DBConn) Recycle() {
	if dbc.conn.IsClosed() {
		dbc.pool.Put(nil)
	} else {
		dbc.pool.Put(dbc)
	}
}

// Kill kills the currently executing query both on MySQL side
// and on the connection side. If no query is executing, it's a no-op.
// Kill will also not kill a query more than once.
func (dbc *DBConn) Kill() {
	killStats.Add("Queries", 1)
	log.Infof("killing query %s", dbc.Current())
	killConn, err := dbc.pool.dbaPool.Get(0)
	if err != nil {
		log.Warningf("Failed to get conn from dba pool: %v", err)
		return
	}
	defer killConn.Recycle()
	sql := fmt.Sprintf("kill %d", dbc.conn.ID())
	_, err = killConn.ExecuteFetch(sql, 10000, false)
	if err != nil {
		log.Errorf("Could not kill query %s: %v", dbc.Current(), err)
	}
}

// Current returns the currently executing query.
func (dbc *DBConn) Current() string {
	return dbc.current.Get()
}

// ID returns the connection id.
func (dbc *DBConn) ID() int64 {
	return dbc.conn.ID()
}

func (dbc *DBConn) reconnect() error {
	dbc.conn.Close()
	newConn, err := dbconnpool.NewDBConnection(dbc.info, mysqlStats)
	if err != nil {
		return err
	}
	dbc.conn = newConn
	return nil
}

func (dbc *DBConn) setDeadline(ctx context.Context) (done chan bool, err error) {
	if ctx.Done() == nil {
		return nil, nil
	}
	done = make(chan bool)
	go func() {
		select {
		case <-ctx.Done():
			// There is a possibility that the query returned very fast,
			// which will cause ctx to get canceled. Check for this condition.
			select {
			case <-done:
				return
			default:
			}
			dbc.Kill()
		case <-done:
			return
		}

		// Verify the query got killed.
		tmr2 := time.NewTimer(15 * time.Second)
		defer tmr2.Stop()
		select {
		case <-tmr2.C:
			internalErrors.Add("HungQuery", 1)
			log.Warningf("Query may be hung: %s", dbc.Current())
		case <-done:
			return
		}
		<-done
		log.Warningf("Hung query returned")
	}()
	return done, nil
}
<commit_msg>remove error as one of return values from DBConn.setDeadline<commit_after>// Copyright 2015, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package tabletserver

import (
	"fmt"
	"time"

	log "github.com/golang/glog"
	mproto "github.com/youtube/vitess/go/mysql/proto"
	"github.com/youtube/vitess/go/sqldb"
	"github.com/youtube/vitess/go/sync2"
	"github.com/youtube/vitess/go/vt/dbconnpool"
	"golang.org/x/net/context"
)

// DBConn is a db connection for tabletserver.
// It performs automatic reconnects as needed.
// Its Execute function has a timeout that can kill
// its own queries and the underlying connection.
// It will also trigger a CheckMySQL whenever applicable.
type DBConn struct {
	conn *dbconnpool.DBConnection
	info *sqldb.ConnParams
	pool *ConnPool

	current sync2.AtomicString
}

// NewDBConn creates a new DBConn. It triggers a CheckMySQL if creation fails.
func NewDBConn(cp *ConnPool, appParams, dbaParams *sqldb.ConnParams) (*DBConn, error) {
	c, err := dbconnpool.NewDBConnection(appParams, mysqlStats)
	if err != nil {
		go checkMySQL()
		return nil, err
	}
	return &DBConn{
		conn: c,
		info: appParams,
		pool: cp,
	}, nil
}

// Exec executes the specified query. If there is a connection error, it will reconnect
// and retry. A failed reconnect will trigger a CheckMySQL.
func (dbc *DBConn) Exec(ctx context.Context, query string, maxrows int, wantfields bool) (*mproto.QueryResult, error) {
	for attempt := 1; attempt <= 2; attempt++ {
		r, err := dbc.execOnce(ctx, query, maxrows, wantfields)
		switch {
		case err == nil:
			return r, nil
		case !IsConnErr(err):
			return nil, NewTabletErrorSql(ErrFail, err)
		case attempt == 2:
			return nil, NewTabletErrorSql(ErrFatal, err)
		}
		err2 := dbc.reconnect()
		if err2 != nil {
			go checkMySQL()
			return nil, NewTabletErrorSql(ErrFatal, err)
		}
	}
	panic("unreachable")
}

func (dbc *DBConn) execOnce(ctx context.Context, query string, maxrows int, wantfields bool) (*mproto.QueryResult, error) {
	dbc.current.Set(query)
	defer dbc.current.Set("")

	done := dbc.setDeadline(ctx)
	if done != nil {
		defer close(done)
	}
	// Uncomment this line for manual testing.
	// defer time.Sleep(20 * time.Second)
	return dbc.conn.ExecuteFetch(query, maxrows, wantfields)
}

// ExecOnce executes the specified query, but does not retry on connection errors.
func (dbc *DBConn) ExecOnce(ctx context.Context, query string, maxrows int, wantfields bool) (*mproto.QueryResult, error) {
	return dbc.execOnce(ctx, query, maxrows, wantfields)
}

// Stream executes the query and streams the results.
func (dbc *DBConn) Stream(ctx context.Context, query string, callback func(*mproto.QueryResult) error, streamBufferSize int) error {
	dbc.current.Set(query)
	defer dbc.current.Set("")

	done := dbc.setDeadline(ctx)
	if done != nil {
		defer close(done)
	}
	return dbc.conn.ExecuteStreamFetch(query, callback, streamBufferSize)
}

// VerifyStrict returns true if MySQL is in STRICT mode.
func (dbc *DBConn) VerifyStrict() bool {
	return dbc.conn.VerifyStrict()
}

// Close closes the DBConn.
func (dbc *DBConn) Close() {
	dbc.conn.Close()
}

// IsClosed returns true if DBConn is closed.
func (dbc *DBConn) IsClosed() bool {
	return dbc.conn.IsClosed()
}

// Recycle returns the DBConn to the pool.
func (dbc *DBConn) Recycle() {
	if dbc.conn.IsClosed() {
		dbc.pool.Put(nil)
	} else {
		dbc.pool.Put(dbc)
	}
}

// Kill kills the currently executing query both on MySQL side
// and on the connection side. If no query is executing, it's a no-op.
// Kill will also not kill a query more than once.
func (dbc *DBConn) Kill() {
	killStats.Add("Queries", 1)
	log.Infof("killing query %s", dbc.Current())
	killConn, err := dbc.pool.dbaPool.Get(0)
	if err != nil {
		log.Warningf("Failed to get conn from dba pool: %v", err)
		return
	}
	defer killConn.Recycle()
	sql := fmt.Sprintf("kill %d", dbc.conn.ID())
	_, err = killConn.ExecuteFetch(sql, 10000, false)
	if err != nil {
		log.Errorf("Could not kill query %s: %v", dbc.Current(), err)
	}
}

// Current returns the currently executing query.
func (dbc *DBConn) Current() string {
	return dbc.current.Get()
}

// ID returns the connection id.
func (dbc *DBConn) ID() int64 {
	return dbc.conn.ID()
}

func (dbc *DBConn) reconnect() error {
	dbc.conn.Close()
	newConn, err := dbconnpool.NewDBConnection(dbc.info, mysqlStats)
	if err != nil {
		return err
	}
	dbc.conn = newConn
	return nil
}

func (dbc *DBConn) setDeadline(ctx context.Context) chan bool {
	if ctx.Done() == nil {
		return nil
	}
	done := make(chan bool)
	go func() {
		select {
		case <-ctx.Done():
			// There is a possibility that the query returned very fast,
			// which will cause ctx to get canceled. Check for this condition.
			select {
			case <-done:
				return
			default:
			}
			dbc.Kill()
		case <-done:
			return
		}

		// Verify the query got killed.
		tmr2 := time.NewTimer(15 * time.Second)
		defer tmr2.Stop()
		select {
		case <-tmr2.C:
			internalErrors.Add("HungQuery", 1)
			log.Warningf("Query may be hung: %s", dbc.Current())
		case <-done:
			return
		}
		<-done
		log.Warningf("Hung query returned")
	}()
	return done
}
<|endoftext|>


<commit_before>/*
Copyright 2017 Google Inc.

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 vtexplain analyzes a set of sql statements and returns the
// corresponding vtgate and vttablet query plans that will be executed
// on the given statements
package vtexplain

import (
	"bytes"
	"fmt"
	"strings"

	log "github.com/golang/glog"

	"github.com/youtube/vitess/go/jsonutil"
	"github.com/youtube/vitess/go/vt/sqlparser"
	"github.com/youtube/vitess/go/vt/vtgate/engine"

	querypb "github.com/youtube/vitess/go/vt/proto/query"
)

// Options to control the explain process
type Options struct {
	// NumShards indicates the number of shards in the topology
	NumShards int

	// ReplicationMode must be set to either "ROW" or "STATEMENT" before
	// initialization
	ReplicationMode string

	// Normalize controls whether or not vtgate does query normalization
	Normalize bool
}

// TabletQuery defines a query that was sent to a given tablet and how it was
// processed in mysql
type TabletQuery struct {
	// SQL command sent to the given tablet
	SQL string

	// BindVars sent with the command
	BindVars map[string]*querypb.BindVariable

	// The actual queries executed by mysql
	MysqlQueries []string
}

// MarshalJSON renders the json structure
func (tq *TabletQuery) MarshalJSON() ([]byte, error) {
	// Convert Bindvars to strings for nicer output
	bindVars := make(map[string]string)
	for k, v := range tq.BindVars {
		var b bytes.Buffer
		sqlparser.EncodeValue(&b, v)
		bindVars[k] = b.String()
	}

	return jsonutil.MarshalNoEscape(&struct {
		SQL          string
		BindVars     map[string]string
		MysqlQueries []string
	}{
		SQL:          tq.SQL,
		BindVars:     bindVars,
		MysqlQueries: tq.MysqlQueries,
	})
}

// Plan defines how vitess will execute a given sql query, including the vtgate
// query plans and all queries run on each tablet.
type Plan struct {
	// original sql statement
	SQL string

	// the vtgate plan(s)
	Plans []*engine.Plan

	// list of queries / bind vars sent to each tablet
	TabletQueries map[string][]*TabletQuery
}

const (
	vtexplainCell = "explainCell"
)

// Init sets up the fake execution environment
func Init(vSchemaStr, sqlSchema string, opts *Options) error {
	// Verify options
	if opts.ReplicationMode != "ROW" && opts.ReplicationMode != "STATEMENT" {
		return fmt.Errorf("invalid replication mode \"%s\"", opts.ReplicationMode)
	}

	err := initVtgateExecutor(vSchemaStr, opts)
	if err != nil {
		return fmt.Errorf("initVtgateExecutor: %v", err)
	}

	parsedDDLs, err := parseSchema(sqlSchema)
	if err != nil {
		return fmt.Errorf("parseSchema: %v", err)
	}

	err = initTabletEnvironment(parsedDDLs, opts)
	if err != nil {
		return fmt.Errorf("initTabletEnvironment: %v", err)
	}

	return nil
}

func parseSchema(sqlSchema string) ([]*sqlparser.DDL, error) {
	parsedDDLs := make([]*sqlparser.DDL, 0, 16)
	for _, sql := range strings.Split(sqlSchema, ";") {
		s := sqlparser.StripLeadingComments(sql)
		s, _ = sqlparser.SplitTrailingComments(sql)
		s = strings.TrimSpace(s)
		if s == "" {
			continue
		}

		stmt, err := sqlparser.Parse(sql)
		if err != nil {
			log.Errorf("ERROR: failed to parse sql: %s, got error: %v", sql, err)
			continue
		}
		ddl, ok := stmt.(*sqlparser.DDL)
		if !ok {
			log.Infof("ignoring non-DDL statement: %s", sql)
			continue
		}
		if ddl.Action != sqlparser.CreateStr {
			log.Infof("ignoring %s table statement", ddl.Action)
			continue
		}
		if ddl.TableSpec == nil {
			log.Errorf("invalid create table statement: %s", sql)
			continue
		}
		parsedDDLs = append(parsedDDLs, ddl)
	}
	return parsedDDLs, nil
}

// Run the explain analysis on the given queries
func Run(sqlStr string) ([]*Plan, error) {
	plans := make([]*Plan, 0, 16)

	for _, sql := range strings.Split(sqlStr, ";") {
		s := strings.TrimSpace(sql)
		if s != "" {
			plan, err := getPlan(s)
			if err != nil {
				return nil, err
			}
			plans = append(plans, plan)
		}
	}

	return plans, nil
}

func getPlan(sql string) (*Plan, error) {
	plans, tabletQueries, err := vtgateExecute(sql)
	if err != nil {
		return nil, err
	}
	for _, tqs := range tabletQueries {
		for _, tq := range tqs {
			mqs, err := fakeTabletExecute(tq.SQL, tq.BindVars)
			if err != nil {
				return nil, fmt.Errorf("fakeTabletExecute: %v", err)
			}
			tq.MysqlQueries = mqs
		}

	}

	return &Plan{
		SQL:           sql,
		Plans:         plans,
		TabletQueries: tabletQueries,
	}, nil
}
<commit_msg>rework vtexplain's Run function to better handle comments<commit_after>/*
Copyright 2017 Google Inc.

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 vtexplain analyzes a set of sql statements and returns the
// corresponding vtgate and vttablet query plans that will be executed
// on the given statements
package vtexplain

import (
	"bytes"
	"fmt"
	"strings"

	log "github.com/golang/glog"

	"github.com/youtube/vitess/go/jsonutil"
	"github.com/youtube/vitess/go/vt/sqlparser"
	"github.com/youtube/vitess/go/vt/vtgate/engine"

	querypb "github.com/youtube/vitess/go/vt/proto/query"
)

// Options to control the explain process
type Options struct {
	// NumShards indicates the number of shards in the topology
	NumShards int

	// ReplicationMode must be set to either "ROW" or "STATEMENT" before
	// initialization
	ReplicationMode string

	// Normalize controls whether or not vtgate does query normalization
	Normalize bool
}

// TabletQuery defines a query that was sent to a given tablet and how it was
// processed in mysql
type TabletQuery struct {
	// SQL command sent to the given tablet
	SQL string

	// BindVars sent with the command
	BindVars map[string]*querypb.BindVariable

	// The actual queries executed by mysql
	MysqlQueries []string
}

// MarshalJSON renders the json structure
func (tq *TabletQuery) MarshalJSON() ([]byte, error) {
	// Convert Bindvars to strings for nicer output
	bindVars := make(map[string]string)
	for k, v := range tq.BindVars {
		var b bytes.Buffer
		sqlparser.EncodeValue(&b, v)
		bindVars[k] = b.String()
	}

	return jsonutil.MarshalNoEscape(&struct {
		SQL          string
		BindVars     map[string]string
		MysqlQueries []string
	}{
		SQL:          tq.SQL,
		BindVars:     bindVars,
		MysqlQueries: tq.MysqlQueries,
	})
}

// Plan defines how vitess will execute a given sql query, including the vtgate
// query plans and all queries run on each tablet.
type Plan struct {
	// original sql statement
	SQL string

	// the vtgate plan(s)
	Plans []*engine.Plan

	// list of queries / bind vars sent to each tablet
	TabletQueries map[string][]*TabletQuery
}

const (
	vtexplainCell = "explainCell"
)

// Init sets up the fake execution environment
func Init(vSchemaStr, sqlSchema string, opts *Options) error {
	// Verify options
	if opts.ReplicationMode != "ROW" && opts.ReplicationMode != "STATEMENT" {
		return fmt.Errorf("invalid replication mode \"%s\"", opts.ReplicationMode)
	}

	err := initVtgateExecutor(vSchemaStr, opts)
	if err != nil {
		return fmt.Errorf("initVtgateExecutor: %v", err)
	}

	parsedDDLs, err := parseSchema(sqlSchema)
	if err != nil {
		return fmt.Errorf("parseSchema: %v", err)
	}

	err = initTabletEnvironment(parsedDDLs, opts)
	if err != nil {
		return fmt.Errorf("initTabletEnvironment: %v", err)
	}

	return nil
}

func parseSchema(sqlSchema string) ([]*sqlparser.DDL, error) {
	parsedDDLs := make([]*sqlparser.DDL, 0, 16)
	for _, sql := range strings.Split(sqlSchema, ";") {
		s := sqlparser.StripLeadingComments(sql)
		s, _ = sqlparser.SplitTrailingComments(sql)
		s = strings.TrimSpace(s)
		if s == "" {
			continue
		}

		stmt, err := sqlparser.Parse(sql)
		if err != nil {
			log.Errorf("ERROR: failed to parse sql: %s, got error: %v", sql, err)
			continue
		}
		ddl, ok := stmt.(*sqlparser.DDL)
		if !ok {
			log.Infof("ignoring non-DDL statement: %s", sql)
			continue
		}
		if ddl.Action != sqlparser.CreateStr {
			log.Infof("ignoring %s table statement", ddl.Action)
			continue
		}
		if ddl.TableSpec == nil {
			log.Errorf("invalid create table statement: %s", sql)
			continue
		}
		parsedDDLs = append(parsedDDLs, ddl)
	}
	return parsedDDLs, nil
}

// Run the explain analysis on the given queries
func Run(sql string) ([]*Plan, error) {
	plans := make([]*Plan, 0, 16)

	for {
		// Need to strip comments in a loop to handle multiple comments
		// in a row.
		for {
			s := sqlparser.StripLeadingComments(sql)
			if s == sql {
				break
			}
			sql = s
		}
		rem := ""
		idx := strings.Index(sql, ";")
		if idx != -1 {
			rem = sql[idx+1:]
			sql = sql[:idx]
		}

		if sql != "" {
			plan, err := getPlan(sql)
			if err != nil {
				return nil, err
			}
			plans = append(plans, plan)
		}

		sql = rem
		if sql == "" {
			break
		}
	}

	return plans, nil
}

func getPlan(sql string) (*Plan, error) {
	plans, tabletQueries, err := vtgateExecute(sql)
	if err != nil {
		return nil, err
	}
	for _, tqs := range tabletQueries {
		for _, tq := range tqs {
			mqs, err := fakeTabletExecute(tq.SQL, tq.BindVars)
			if err != nil {
				return nil, fmt.Errorf("fakeTabletExecute: %v", err)
			}
			tq.MysqlQueries = mqs
		}

	}

	return &Plan{
		SQL:           sql,
		Plans:         plans,
		TabletQueries: tabletQueries,
	}, nil
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/sha512"
	"encoding/base64"
	"fmt"
	"io"
	"log"
	"net/http"
	"strings"
)

// during alpha, we use basic auth for alpha access
// basic auth is chosen because it does not interfere with any other techniques used by this project
// the alphaUsers map holds all alpha users and their sha512 hashed passwords.
// the password is to be salted with the alphaSalt (e.g. <password>+alphaSalt)
// feel free to add yourself or other project members to the map
var alphaSalt = "che9ohlu0xie9yaW6aeg"
var alphaUsers = map[string]string{
	"GeertJohan": "8e764667137fc73b16a3a3cd43a6a9314c8c2214215306f563068bd91b7f5de06da19815b10e3b6719c7f9cd4499b26e13738fadba0b244c4055f1d7af0100b8",
}

func alphaCheckBasicAuth(r *http.Request) bool {
	// retrieve auth header
	authHeader := r.Header.Get("Authorization")
	if len(authHeader) < 7 {
		// no or invalid auth given
		return false
	}
	if authHeader[:6] != "Basic " {
		// invalid auth type
		return false
	}

	// decode auth data
	authData, err := base64.StdEncoding.DecodeString(authHeader[6:])
	if err != nil {
		log.Printf("Could not decode auth. %s\n", err)
		return false
	}
	authDataSlice := strings.SplitN(string(authData), ":", 2)
	givenUsername := authDataSlice[0]
	givenPassword := authDataSlice[1]

	// retrieved hashed password from alphaUsers map
	correctPasswordHashed, userExists := alphaUsers[givenUsername]
	if !userExists {
		// user does not exist
		return false
	}
	log.Println("user exists. hash is: ")
	log.Println(correctPasswordHashed)

	// hash retrieved password, format as hex string
	passwordHasher := sha512.New()
	io.WriteString(passwordHasher, givenPassword)
	io.WriteString(passwordHasher, alphaSalt)
	givenPasswordHashed := fmt.Sprintf("%x", passwordHasher.Sum(nil))
	log.Println(givenPasswordHashed)
	// password matches?
	if givenPasswordHashed == correctPasswordHashed {
		// yay! correct auth!
		log.Println("correct auth")
		return true
	}

	return false
}
<commit_msg>Removed logging from alpha.go Added extra check on base auth correctness<commit_after>package main

import (
	"crypto/sha512"
	"encoding/base64"
	"fmt"
	"io"
	"net/http"
	"strings"
)

// during alpha, we use basic auth for alpha access
// basic auth is chosen because it does not interfere with any other techniques used by this project
// the alphaUsers map holds all alpha users and their sha512 hashed passwords.
// the password is to be salted with the alphaSalt (e.g. <password>+alphaSalt)
// feel free to add yourself or other project members to the map
var alphaSalt = "che9ohlu0xie9yaW6aeg"
var alphaUsers = map[string]string{
	"GeertJohan": "8e764667137fc73b16a3a3cd43a6a9314c8c2214215306f563068bd91b7f5de06da19815b10e3b6719c7f9cd4499b26e13738fadba0b244c4055f1d7af0100b8",
}

func alphaCheckBasicAuth(r *http.Request) bool {
	// retrieve auth header
	authHeader := r.Header.Get("Authorization")
	if len(authHeader) < 7 {
		// no or invalid auth given
		return false
	}
	if authHeader[:6] != "Basic " {
		// invalid auth type
		return false
	}

	// decode auth data
	authData, err := base64.StdEncoding.DecodeString(authHeader[6:])
	if err != nil {
		return false
	}
	authDataSlice := strings.SplitN(string(authData), ":", 2)
	if len(authDataSlice) != 2 {
		return false
	}
	givenUsername := authDataSlice[0]
	givenPassword := authDataSlice[1]

	// retrieved hashed password from alphaUsers map
	correctPasswordHashed, userExists := alphaUsers[givenUsername]
	if !userExists {
		// user does not exist
		return false
	}

	// hash retrieved password, format as hex string
	passwordHasher := sha512.New()
	io.WriteString(passwordHasher, givenPassword)
	io.WriteString(passwordHasher, alphaSalt)
	givenPasswordHashed := fmt.Sprintf("%x", passwordHasher.Sum(nil))
	// password matches?
	if givenPasswordHashed == correctPasswordHashed {
		// yay! correct auth!
		return true
	}

	return false
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"sort"
	"text/tabwriter"
	"time"

	"github.com/krasoffski/gomill/memosort"
)

// Track represent information about music track.
type Track struct {
	Title  string
	Artist string
	Album  string
	Year   int
	Length time.Duration
}

func printTracks(tracks []*Track) {
	const format = "%v\t%v\t%v\t%v\t%v\t\n"
	tw := new(tabwriter.Writer).Init(os.Stdout, 0, 8, 2, ' ', 0)
	fmt.Fprintf(tw, format, "Title", "Artist", "Album", "Year", "Length")
	fmt.Fprintf(tw, format, "-----", "------", "-----", "----", "------")
	for _, t := range tracks {
		fmt.Fprintf(tw, format, t.Title, t.Artist, t.Album, t.Year, t.Length)
	}
	tw.Flush()
}

var tracks = []*Track{
	{"Go", "Delilah", "From the Roots Up", 2012, length("3m38s")},
	{"Go", "Moby", "Moby", 1992, length("3m37s")},
	{"Go Ahead", "Alicia Keys", "As I Am", 2007, length("4m36s")},
	{"Ready 2 Go", "Martin Solveig", "Smash", 2011, length("4m24s")},
}

func length(s string) time.Duration {
	d, err := time.ParseDuration(s)
	if err != nil {
		panic(s)
	}
	return d
}

func main() {
	m := memosort.New()
	sort.Slice(tracks, m.By(func(i, j int) bool {
		return tracks[i].Title < tracks[j].Title
	}))
	sort.Slice(tracks, m.By(func(i, j int) bool {
		return tracks[i].Year < tracks[j].Year
	}))
	sort.Slice(tracks, m.By(func(i, j int) bool {
		return tracks[i].Length < tracks[j].Length
	}))

	printTracks(tracks)
}
<commit_msg>[7.9] Added draft solution.<commit_after>package main

import (
	"fmt"
	"html/template"
	"log"
	"net/http"
	"os"
	"sort"
	"text/tabwriter"
	"time"

	"github.com/krasoffski/gomill/memosort"
)

// Track represent information about music track.
type Track struct {
	Title  string
	Artist string
	Album  string
	Year   int
	Length time.Duration
}

func printTracks(tracks []*Track) {
	const format = "%v\t%v\t%v\t%v\t%v\t\n"
	tw := new(tabwriter.Writer).Init(os.Stdout, 0, 8, 2, ' ', 0)
	fmt.Fprintf(tw, format, "Title", "Artist", "Album", "Year", "Length")
	fmt.Fprintf(tw, format, "-----", "------", "-----", "----", "------")
	for _, t := range tracks {
		fmt.Fprintf(tw, format, t.Title, t.Artist, t.Album, t.Year, t.Length)
	}
	tw.Flush()
}

var tracks = []*Track{
	{"Go", "Delilah", "From the Roots Up", 2012, length("3m38s")},
	{"Go", "Moby", "Moby", 1992, length("3m37s")},
	{"Go Ahead", "Alicia Keys", "As I Am", 2007, length("4m36s")},
	{"Ready 2 Go", "Martin Solveig", "Smash", 2011, length("4m24s")},
}

func length(s string) time.Duration {
	d, err := time.ParseDuration(s)
	if err != nil {
		panic(s)
	}
	return d
}

var tracksTemplate = template.Must(template.New("tracksTemplate").Parse(`

	Title
	Artist
	Album
	Year
	Length

{{range .}}

	{{.Title}}
	{{.Artist}}
	{{.Album}}
	{{.Year}}
	{{.Length}}

{{end}}
`))

func handler(w http.ResponseWriter, r *http.Request) {
	switch r.URL.Query().Get("sort") {
	case "title":
		sort.Slice(tracks,
			func(i, j int) bool { return tracks[i].Title < tracks[j].Title })
	case "artist":
		sort.Slice(tracks,
			func(i, j int) bool { return tracks[i].Artist < tracks[j].Artist })
	case "album":
		sort.Slice(tracks,
			func(i, j int) bool { return tracks[i].Album < tracks[j].Album })
	case "year":
		sort.Slice(tracks,
			func(i, j int) bool { return tracks[i].Year < tracks[j].Year })
	case "length":
		sort.Slice(tracks,
			func(i, j int) bool { return tracks[i].Length < tracks[j].Length })
	}

	if err := tracksTemplate.Execute(w, tracks); err != nil {
		log.Print(err)
	}
}

func main() {
	// TODO: remove duplicated code for less functions.
	m := memosort.New()
	sort.Slice(tracks, m.By(func(i, j int) bool {
		return tracks[i].Title < tracks[j].Title
	}))
	sort.Slice(tracks, m.By(func(i, j int) bool {
		return tracks[i].Year < tracks[j].Year
	}))
	sort.Slice(tracks, m.By(func(i, j int) bool {
		return tracks[i].Length < tracks[j].Length
	}))
	printTracks(tracks)
	fmt.Println("Starting server...")
	http.HandleFunc("/", handler)
	log.Fatal(http.ListenAndServe(":8000", nil))
}
<|endoftext|>


<commit_before>/*
Licensed to the Apache Software Foundation (ASF) under one
or more contributor license agreements.  See the NOTICE file
distributed with this work for additional information
regarding copyright ownership.  The ASF licenses this file
to you 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 gremlingo

import "github.com/google/uuid"

// request represents a request to the server.
type request struct {
	requestID uuid.UUID
	op        string
	processor string
	args      map[string]interface{}
}

const sessionProcessor = "session"

const stringOp = "eval"
const stringProcessor = ""

// Bindings should be a key-object map (different from Binding class in Bytecode).
func makeStringRequest(stringGremlin string, traversalSource string, sessionId string, bindings ...map[string]interface{}) (req request) {
	newProcessor := stringProcessor
	newArgs := map[string]interface{}{
		"gremlin": stringGremlin,
		"aliases": map[string]interface{}{
			"g": traversalSource,
		},
	}
	if sessionId != "" {
		newProcessor = sessionProcessor
		newArgs["session"] = sessionId
	}
	if len(bindings) > 0 {
		newArgs["bindings"] = bindings[0]
	}
	return request{
		requestID: uuid.New(),
		op:        stringOp,
		processor: newProcessor,
		args:      newArgs,
	}
}

const bytecodeOp = "bytecode"
const bytecodeProcessor = "traversal"
const authOp = "authentication"
const authProcessor = "traversal"

func makeBytecodeRequest(bytecodeGremlin *Bytecode, traversalSource string, sessionId string) (req request) {
	newProcessor := bytecodeProcessor
	newArgs := map[string]interface{}{
		"gremlin": *bytecodeGremlin,
		"aliases": map[string]interface{}{
			"g": traversalSource,
		},
	}
	if sessionId != "" {
		newProcessor = sessionProcessor
		newArgs["session"] = sessionId
	}

	for k, v := range extractReqArgs(bytecodeGremlin) {
		newArgs[k] = v
	}

	return request{
		requestID: uuid.New(),
		op:        bytecodeOp,
		processor: newProcessor,
		args:      newArgs,
	}
}

// allowedReqArgs contains the arguments that will be extracted from the
// bytecode and sent with the request.
var allowedReqArgs = map[string]bool{
	"evaluationTimeout": true,
	"batchSize":         true,
	"requestId":         true,
	"userAgent":         true,
}

// extractReqArgs extracts request arguments from the provided bytecode.
func extractReqArgs(bytecode *Bytecode) map[string]interface{} {
	args := make(map[string]interface{})

	for _, insn := range bytecode.sourceInstructions {
		switch insn.operator {
		case "withStrategies":
			for k, v := range extractWithStrategiesReqArgs(insn) {
				args[k] = v
			}
		case "with":
			if k, v := extractWithReqArgs(insn); k != "" {
				args[k] = v
			}
		}
	}

	return args
}

// extractWithStrategiesReqArgs extracts request arguments from the passed
// "withStrategies" source instruction. Only OptionsStrategy is considered.
func extractWithStrategiesReqArgs(insn instruction) map[string]interface{} {
	args := make(map[string]interface{})

	for _, strategyInterface := range insn.arguments {
		strategy, ok := strategyInterface.(*traversalStrategy)
		if !ok {
			// (*GraphTraversalSource).WithStrategies accepts
			// TraversalStrategy parameters only. Thus, this
			// should be unreachable.
			continue
		}

		if strategy.name != decorationNamespace+"OptionsStrategy" {
			continue
		}

		for k, v := range strategy.configuration {
			if allowedReqArgs[k] {
				args[k] = v
			}
		}
	}

	return args
}

// extractWithReqArgs extracts a request argument from the passed "with" source
// instruction.
func extractWithReqArgs(insn instruction) (key string, value interface{}) {
	if len(insn.arguments) != 2 {
		// (*GraphTraversalSource).With accepts two parameters. Thus,
		// this should be unreachable.
		return "", nil
	}

	key, ok := insn.arguments[0].(string)
	if !ok {
		return "", nil
	}

	if !allowedReqArgs[key] {
		return "", nil
	}

	return key, insn.arguments[1]
}

func makeBasicAuthRequest(auth string) (req request) {
	return request{
		requestID: uuid.New(),
		op:        authOp,
		processor: authProcessor,
		args: map[string]interface{}{
			"sasl": auth,
		},
	}
}

func makeCloseSessionRequest(sessionId string) request {
	return request{
		requestID: uuid.New(),
		op:        "close",
		processor: "session",
		args: map[string]interface{}{
			"session": sessionId,
		},
	}
}
<commit_msg>gremlin-go: rename function "extractWithReqArgs" to "extractWithReqArg"<commit_after>/*
Licensed to the Apache Software Foundation (ASF) under one
or more contributor license agreements.  See the NOTICE file
distributed with this work for additional information
regarding copyright ownership.  The ASF licenses this file
to you 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 gremlingo

import "github.com/google/uuid"

// request represents a request to the server.
type request struct {
	requestID uuid.UUID
	op        string
	processor string
	args      map[string]interface{}
}

const sessionProcessor = "session"

const stringOp = "eval"
const stringProcessor = ""

// Bindings should be a key-object map (different from Binding class in Bytecode).
func makeStringRequest(stringGremlin string, traversalSource string, sessionId string, bindings ...map[string]interface{}) (req request) {
	newProcessor := stringProcessor
	newArgs := map[string]interface{}{
		"gremlin": stringGremlin,
		"aliases": map[string]interface{}{
			"g": traversalSource,
		},
	}
	if sessionId != "" {
		newProcessor = sessionProcessor
		newArgs["session"] = sessionId
	}
	if len(bindings) > 0 {
		newArgs["bindings"] = bindings[0]
	}
	return request{
		requestID: uuid.New(),
		op:        stringOp,
		processor: newProcessor,
		args:      newArgs,
	}
}

const bytecodeOp = "bytecode"
const bytecodeProcessor = "traversal"
const authOp = "authentication"
const authProcessor = "traversal"

func makeBytecodeRequest(bytecodeGremlin *Bytecode, traversalSource string, sessionId string) (req request) {
	newProcessor := bytecodeProcessor
	newArgs := map[string]interface{}{
		"gremlin": *bytecodeGremlin,
		"aliases": map[string]interface{}{
			"g": traversalSource,
		},
	}
	if sessionId != "" {
		newProcessor = sessionProcessor
		newArgs["session"] = sessionId
	}

	for k, v := range extractReqArgs(bytecodeGremlin) {
		newArgs[k] = v
	}

	return request{
		requestID: uuid.New(),
		op:        bytecodeOp,
		processor: newProcessor,
		args:      newArgs,
	}
}

// allowedReqArgs contains the arguments that will be extracted from the
// bytecode and sent with the request.
var allowedReqArgs = map[string]bool{
	"evaluationTimeout": true,
	"batchSize":         true,
	"requestId":         true,
	"userAgent":         true,
}

// extractReqArgs extracts request arguments from the provided bytecode.
func extractReqArgs(bytecode *Bytecode) map[string]interface{} {
	args := make(map[string]interface{})

	for _, insn := range bytecode.sourceInstructions {
		switch insn.operator {
		case "withStrategies":
			for k, v := range extractWithStrategiesReqArgs(insn) {
				args[k] = v
			}
		case "with":
			if k, v := extractWithReqArg(insn); k != "" {
				args[k] = v
			}
		}
	}

	return args
}

// extractWithStrategiesReqArgs extracts request arguments from the passed
// "withStrategies" source instruction. Only OptionsStrategy is considered.
func extractWithStrategiesReqArgs(insn instruction) map[string]interface{} {
	args := make(map[string]interface{})

	for _, strategyInterface := range insn.arguments {
		strategy, ok := strategyInterface.(*traversalStrategy)
		if !ok {
			// (*GraphTraversalSource).WithStrategies accepts
			// TraversalStrategy parameters only. Thus, this
			// should be unreachable.
			continue
		}

		if strategy.name != decorationNamespace+"OptionsStrategy" {
			continue
		}

		for k, v := range strategy.configuration {
			if allowedReqArgs[k] {
				args[k] = v
			}
		}
	}

	return args
}

// extractWithReqArg extracts a request argument from the passed "with" source
// instruction.
func extractWithReqArg(insn instruction) (key string, value interface{}) {
	if len(insn.arguments) != 2 {
		// (*GraphTraversalSource).With accepts two parameters. Thus,
		// this should be unreachable.
		return "", nil
	}

	key, ok := insn.arguments[0].(string)
	if !ok {
		return "", nil
	}

	if !allowedReqArgs[key] {
		return "", nil
	}

	return key, insn.arguments[1]
}

func makeBasicAuthRequest(auth string) (req request) {
	return request{
		requestID: uuid.New(),
		op:        authOp,
		processor: authProcessor,
		args: map[string]interface{}{
			"sasl": auth,
		},
	}
}

func makeCloseSessionRequest(sessionId string) request {
	return request{
		requestID: uuid.New(),
		op:        "close",
		processor: "session",
		args: map[string]interface{}{
			"session": sessionId,
		},
	}
}
<|endoftext|>


<commit_before>package algoliaconnector

import (
	"koding/db/mongodb/modelhelper"
	"math/rand"
	"socialapi/config"
	"socialapi/models"
	"strconv"
	"testing"
	"time"

	. "github.com/smartystreets/goconvey/convey"
)

func TestChannelCreated(t *testing.T) {
	runner, handler := getTestHandler()
	defer runner.Close()

	Convey("given some fake topic channel", t, func() {
		mockTopic := models.NewChannel()
		mockTopic.TypeConstant = models.Channel_TYPE_TOPIC
		Convey("it should save the document to algolia", func() {
			err := handler.ChannelCreated(mockTopic)
			So(err, ShouldBeNil)

			err = makeSureChannel(handler, mockTopic.Id, func(record map[string]interface{}, err error) bool {
				if err != nil {
					return false
				}

				return true
			})
			So(err, ShouldBeNil)
			rec, err := handler.get(IndexTopics, strconv.FormatInt(mockTopic.Id, 10))
			So(err, ShouldBeNil)
			So(rec["_tags"], ShouldNotBeNil)
			So(len(rec["_tags"].([]interface{})), ShouldBeGreaterThan, 0)
		})
	})

	Convey("given some fake non-topic channel", t, func() {
		mockTopic := models.NewChannel()
		mockTopic.TypeConstant = models.Channel_TYPE_PRIVATE_MESSAGE
		Convey("it should save the document to algolia", func() {
			err := handler.ChannelCreated(mockTopic)
			So(err, ShouldBeNil)
		})
	})
}

func TestChannelUpdated(t *testing.T) {
	runner, handler := getTestHandler()
	defer runner.Close()

	rand.Seed(time.Now().UnixNano())

	appConfig := config.MustRead(runner.Conf.Path)
	modelhelper.Initialize(appConfig.Mongo)
	defer modelhelper.Close()

	Convey("given some fake topic channel", t, func() {
		mockTopic := models.NewChannel()
		mockTopic.Id = rand.Int63()
		mockTopic.TypeConstant = models.Channel_TYPE_TOPIC
		Convey("it should save the document to algolia", func() {
			err := handler.ChannelCreated(mockTopic)
			So(err, ShouldBeNil)
			err = makeSureChannel(handler, mockTopic.Id, func(record map[string]interface{}, err error) bool {
				if err != nil {
					return false
				}

				return true
			})

			So(err, ShouldBeNil)

			Convey("given some existing topic channel", func() {
				mockTopic.TypeConstant = models.Channel_TYPE_LINKED_TOPIC
				Convey("it should be able to remove it", func() {
					err := handler.ChannelUpdated(mockTopic)
					So(err, ShouldBeNil)

					err = makeSureChannel(handler, mockTopic.Id, func(record map[string]interface{}, err error) bool {
						if IsAlgoliaError(err, ErrAlgoliaObjectIdNotFoundMsg) {
							return true
						}

						return false
					})

					So(err, ShouldBeNil)

					Convey("removing a deleted channel should return success", func() {
						err := handler.ChannelUpdated(mockTopic)
						So(err, ShouldBeNil)
					})
				})
			})
			Convey("removing a non-existing channel should return success", func() {
				mockTopic.Id++
				err := handler.ChannelUpdated(mockTopic)
				So(err, ShouldBeNil)
			})
		})
	})
}

// makeSureChannel checks if the given id's get request returns the desired err,
// it will re-try every 100ms until deadline of 2 minutes. Algolia doesnt index
// the records right away, so try to go to a desired state
func makeSureChannel(handler *Controller, id int64, f func(map[string]interface{}, error) bool) error {
	deadLine := time.After(TestTimeout)
	tick := time.Tick(time.Millisecond * 100)
	for {
		select {
		case <-tick:
			record, err := handler.get(IndexTopics, strconv.FormatInt(id, 10))
			if f(record, err) {
				return nil
			}
		case <-deadLine:
			return errDeadline
		}
	}
}
<commit_msg>Socialapi: channel should be created before sending to AlgoliaConnector<commit_after>package algoliaconnector

import (
	"koding/db/mongodb/modelhelper"
	"math/rand"
	"socialapi/config"
	"socialapi/models"
	"strconv"
	"testing"
	"time"

	. "github.com/smartystreets/goconvey/convey"
)

func TestChannelCreated(t *testing.T) {
	runner, handler := getTestHandler()
	defer runner.Close()

	appConfig := config.MustRead(runner.Conf.Path)
	modelhelper.Initialize(appConfig.Mongo)
	defer modelhelper.Close()

	Convey("given some fake topic channel", t, func() {
		groupName := models.RandomGroupName()
		acc1, err := models.CreateAccountInBothDbs()
		So(err, ShouldBeNil)

		// we need group channel,because we are injecting it's id into channel
		// as tag
		models.CreateTypedGroupedChannelWithTest(
			acc1.Id,
			models.Channel_TYPE_GROUP,
			groupName,
		)

		mockTopic := models.CreateTypedGroupedChannelWithTest(
			acc1.Id,
			models.Channel_TYPE_TOPIC,
			groupName,
		)

		Convey("it should save the document to algolia", func() {
			err := handler.ChannelCreated(mockTopic)
			So(err, ShouldBeNil)

			err = makeSureChannel(handler, mockTopic.Id, func(record map[string]interface{}, err error) bool {
				if err != nil {
					return false
				}

				return true
			})
			So(err, ShouldBeNil)
			rec, err := handler.get(IndexTopics, strconv.FormatInt(mockTopic.Id, 10))
			So(err, ShouldBeNil)
			So(rec["_tags"], ShouldNotBeNil)
			So(len(rec["_tags"].([]interface{})), ShouldBeGreaterThan, 0)
		})
	})

	Convey("given some fake non-topic channel", t, func() {
		mockTopic := models.NewChannel()
		mockTopic.TypeConstant = models.Channel_TYPE_PRIVATE_MESSAGE
		Convey("it should save the document to algolia", func() {
			err := handler.ChannelCreated(mockTopic)
			So(err, ShouldBeNil)
		})
	})
}

func TestChannelUpdated(t *testing.T) {
	runner, handler := getTestHandler()
	defer runner.Close()

	rand.Seed(time.Now().UnixNano())

	appConfig := config.MustRead(runner.Conf.Path)
	modelhelper.Initialize(appConfig.Mongo)
	defer modelhelper.Close()

	Convey("given some fake topic channel", t, func() {
		mockTopic := models.NewChannel()
		mockTopic.Id = rand.Int63()
		mockTopic.TypeConstant = models.Channel_TYPE_TOPIC
		Convey("it should save the document to algolia", func() {
			err := handler.ChannelCreated(mockTopic)
			So(err, ShouldBeNil)
			err = makeSureChannel(handler, mockTopic.Id, func(record map[string]interface{}, err error) bool {
				if err != nil {
					return false
				}

				return true
			})

			So(err, ShouldBeNil)

			Convey("given some existing topic channel", func() {
				mockTopic.TypeConstant = models.Channel_TYPE_LINKED_TOPIC
				Convey("it should be able to remove it", func() {
					err := handler.ChannelUpdated(mockTopic)
					So(err, ShouldBeNil)

					err = makeSureChannel(handler, mockTopic.Id, func(record map[string]interface{}, err error) bool {
						if IsAlgoliaError(err, ErrAlgoliaObjectIdNotFoundMsg) {
							return true
						}

						return false
					})

					So(err, ShouldBeNil)

					Convey("removing a deleted channel should return success", func() {
						err := handler.ChannelUpdated(mockTopic)
						So(err, ShouldBeNil)
					})
				})
			})
			Convey("removing a non-existing channel should return success", func() {
				mockTopic.Id++
				err := handler.ChannelUpdated(mockTopic)
				So(err, ShouldBeNil)
			})
		})
	})
}

// makeSureChannel checks if the given id's get request returns the desired err,
// it will re-try every 100ms until deadline of 2 minutes. Algolia doesnt index
// the records right away, so try to go to a desired state
func makeSureChannel(handler *Controller, id int64, f func(map[string]interface{}, error) bool) error {
	deadLine := time.After(TestTimeout)
	tick := time.Tick(time.Millisecond * 100)
	for {
		select {
		case <-tick:
			record, err := handler.get(IndexTopics, strconv.FormatInt(id, 10))
			if f(record, err) {
				return nil
			}
		case <-deadLine:
			return errDeadline
		}
	}
}
<|endoftext|>