// CheckQueryResults checks that the rows returned by a query match the expected
// response.
func (sr *SQLRunner) CheckQueryResults(query string, expected [][]string) {
	res := sr.QueryStr(query)
	if !reflect.DeepEqual(res, expected) {
		file, line, _ := caller.Lookup(1)
		sr.Errorf("%s:%d query '%s': expected:\n%v\ngot:%v\n", file, line, query, expected, res)
	}
}

// SetDebugName sets the debug name associated with the transaction which will
// appear in log files and the web UI. Each transaction starts out with an
// automatically assigned debug name composed of the file and line number where
// the transaction was created.
func (txn *Txn) SetDebugName(name string, depth int) {
	file, line, fun := caller.Lookup(depth + 1)
	if name == "" {
		name = fun
	}
	txn.Proto.Name = file + ":" + strconv.Itoa(line) + " " + name
}

func TestLogBacktraceAt(t *testing.T) {
	setFlags()
	defer logging.swap(logging.newBuffers())
	// The peculiar style of this code simplifies line counting and maintenance of the
	// tracing block below.
	var infoLine string
	setTraceLocation := func(file string, line int, delta int) {
		_, file = filepath.Split(file)
		infoLine = fmt.Sprintf("%s:%d", file, line+delta)
		err := logging.traceLocation.Set(infoLine)
		if err != nil {
			t.Fatal("error setting log_backtrace_at: ", err)
		}
	}
	{
		// Start of tracing block. These lines know about each other's relative position.
		file, line, _ := caller.Lookup(0)
		setTraceLocation(file, line, +2) // Two lines between Caller and Info calls.
		Info(context.Background(), "we want a stack trace here")
		if err := logging.traceLocation.Set(""); err != nil {
			t.Fatal(err)
		}
	}
	numAppearances := strings.Count(contents(Severity_INFO), infoLine)
	if numAppearances < 2 {
		// Need 2 appearances, one in the log header and one in the trace:
		//   log_test.go:281: I0511 16:36:06.952398 02238 log_test.go:280] we want a stack trace here
		//   ...
		//   github.com/clog/glog_test.go:280 (0x41ba91)
		//   ...
		// We could be more precise but that would require knowing the details
		// of the traceback format, which may not be dependable.
		t.Fatal("got no trace back; log is ", contents(Severity_INFO))
	}
}

// Stop signals all live workers to stop and then waits for each to
// confirm it has stopped.
func (s *Stopper) Stop() {
	defer s.Recover()
	defer unregister(s)

	file, line, _ := caller.Lookup(1)
	log.Infof(context.TODO(),
		"stop has been called from %s:%d, stopping or quiescing all running tasks", file, line)
	// Don't bother doing stuff cleanly if we're panicking, that would likely
	// block. Instead, best effort only. This cleans up the stack traces,
	// avoids stalls and helps some tests in `./cli` finish cleanly (where
	// panics happen on purpose).
	if r := recover(); r != nil {
		go s.Quiesce()
		close(s.stopper)
		close(s.stopped)
		s.mu.Lock()
		for _, c := range s.mu.closers {
			go c.Close()
		}
		s.mu.Unlock()
		panic(r)
	}

	s.Quiesce()
	close(s.stopper)
	s.stop.Wait()
	s.mu.Lock()
	defer s.mu.Unlock()
	for _, c := range s.mu.closers {
		c.Close()
	}
	close(s.stopped)
}

// EnsureContext checks whether the given context.Context contains a Span. If
// not, it creates one using the provided Tracer and wraps it in the returned
// Span. The returned closure must be called after the request has been fully
// processed.
func EnsureContext(ctx context.Context, tracer opentracing.Tracer) (context.Context, func()) {
	_, _, funcName := caller.Lookup(1)
	if opentracing.SpanFromContext(ctx) == nil {
		sp := tracer.StartSpan(funcName)
		return opentracing.ContextWithSpan(ctx, sp), sp.Finish
	}
	return ctx, func() {}
}

// addStructured creates a structured log entry to be written to the
// specified facility of the logger.
func addStructured(ctx context.Context, s Severity, depth int, format string, args []interface{}) {
	file, line, _ := caller.Lookup(depth + 1)
	msg := makeMessage(ctx, format, args)
	// makeMessage already added the tags when forming msg, we don't want
	// eventInternal to prepend them again.
	eventInternal(ctx, (s >= Severity_ERROR), false /*withTags*/, "%s:%d %s", file, line, msg)
	logging.outputLogEntry(s, file, line, msg)
}

func defTestCase(expected, expectedReverse int, desired ...sqlbase.ColumnOrderInfo) desiredCase {
	// The line number is used to identify testcases in error messages.
	_, line, _ := caller.Lookup(1)
	return desiredCase{
		line:            line,
		desired:         desired,
		expected:        expected,
		expectedReverse: expectedReverse,
	}
}

// Caller returns filename and line number info for the specified stack
// depths. The info is formated as <file>:<line> and each entry is separated
// for a space.
func Caller(depth ...int) string {
	var sep string
	var buf bytes.Buffer
	for _, d := range depth {
		file, line, _ := caller.Lookup(d + 1)
		fmt.Fprintf(&buf, "%s%s:%d", sep, file, line)
		sep = " "
	}
	return buf.String()
}

// RunTaskWithErr adds one to the count of tasks left to quiesce in the system.
// Any worker which is a "first mover" when starting tasks must call this method
// before starting work on a new task. First movers include goroutines launched
// to do periodic work and the kv/db.go gateway which accepts external client
// requests.
//
// If the system is currently quiescing and function f was not called, returns
// an error indicating this condition. Otherwise, returns whatever f returns.
func (s *Stopper) RunTaskWithErr(f func() error) error {
	file, line, _ := caller.Lookup(1)
	key := taskKey{file, line}
	if !s.runPrelude(key) {
		return errUnavailable
	}
	// Call f.
	defer s.Recover()
	defer s.runPostlude(key)
	return f()
}

// checkRestart checks that there are no errors left to inject.
func checkRestarts(t *testing.T, magicVals *filterVals) {
	magicVals.Lock()
	defer magicVals.Unlock()
	for key, count := range magicVals.restartCounts {
		if count != 0 {
			file, line, _ := caller.Lookup(1)
			t.Errorf("%s:%d: INSERT for \"%s\" still has to be retried %d times",
				file, line, key, count)
		}
	}
	for key, count := range magicVals.abortCounts {
		if count != 0 {
			file, line, _ := caller.Lookup(1)
			t.Errorf("%s:%d: INSERT for \"%s\" still has to be aborted %d times",
				file, line, key, count)
		}
	}
	if t.Failed() {
		t.Fatalf("checking error injection failed")
	}
}

// RunLimitedAsyncTask runs function f in a goroutine, using the given
// channel as a semaphore to limit the number of tasks that are run
// concurrently to the channel's capacity. If wait is true, blocks
// until the semaphore is available in order to push back on callers
// that may be trying to create many tasks. If wait is false, returns
// immediately with an error if the semaphore is not
// available. Returns an error if the Stopper is quiescing, in which
// case the function is not executed.
func (s *Stopper) RunLimitedAsyncTask(
	ctx context.Context, sem chan struct{}, wait bool, f func(context.Context),
) error {
	file, line, _ := caller.Lookup(1)
	key := taskKey{file, line}

	// Wait for permission to run from the semaphore.
	select {
	case sem <- struct{}{}:
	case <-ctx.Done():
		return ctx.Err()
	case <-s.ShouldQuiesce():
		return errUnavailable
	default:
		if !wait {
			return ErrThrottled
		}
		log.Infof(context.TODO(), "stopper throttling task from %s:%d due to semaphore", file, line)
		// Retry the select without the default.
		select {
		case sem <- struct{}{}:
		case <-ctx.Done():
			return ctx.Err()
		case <-s.ShouldQuiesce():
			return errUnavailable
		}
	}

	// Check for canceled context: it's possible to get the semaphore even
	// if the context is canceled.
	select {
	case <-ctx.Done():
		<-sem
		return ctx.Err()
	default:
	}

	if !s.runPrelude(key) {
		<-sem
		return errUnavailable
	}

	ctx, span := tracing.ForkCtxSpan(ctx, fmt.Sprintf("%s:%d", file, line))

	go func() {
		defer s.Recover()
		defer s.runPostlude(key)
		defer func() { <-sem }()
		defer tracing.FinishSpan(span)
		f(ctx)
	}()
	return nil
}

// TempDir creates a directory and a function to clean it up at the end of the
// test. If called directly from a test function, pass 0 for depth (which puts
// the test name in the directory). Otherwise, offset depth appropriately.
func TempDir(t testing.TB, depth int) (string, func()) {
	_, _, name := caller.Lookup(depth + 1)
	dir, err := ioutil.TempDir("", name)
	if err != nil {
		t.Fatal(err)
	}
	cleanup := func() {
		if err := os.RemoveAll(dir); err != nil {
			t.Error(err)
		}
	}
	return dir, cleanup
}

// CheckQueryResults checks that the rows returned by a query match the expected
// response.
func (sr *SQLRunner) CheckQueryResults(query string, expected [][]string) {
	file, line, _ := caller.Lookup(1)
	info := fmt.Sprintf("%s:%d query '%s'", file, line, query)

	rows := sr.Query(query)
	cols, err := rows.Columns()
	if err != nil {
		sr.Error(err)
		return
	}
	if len(expected) > 0 && len(cols) != len(expected[0]) {
		sr.Errorf("%s: wrong number of columns %d", info, len(cols))
		return
	}
	vals := make([]interface{}, len(cols))
	for i := range vals {
		vals[i] = new(interface{})
	}
	i := 0
	for ; rows.Next(); i++ {
		if i >= len(expected) {
			sr.Errorf("%s: expected %d rows, got more", info, len(expected))
			return
		}
		if err := rows.Scan(vals...); err != nil {
			sr.Error(err)
			return
		}
		for j, v := range vals {
			if val := *v.(*interface{}); val != nil {
				var s string
				switch t := val.(type) {
				case []byte:
					s = string(t)
				default:
					s = fmt.Sprint(val)
				}
				if expected[i][j] != s {
					sr.Errorf("%s: expected %v, found %v", info, expected[i][j], s)
				}
			} else if expected[i][j] != "NULL" {
				sr.Errorf("%s: expected %v, found %v", info, expected[i][j], "NULL")
			}
		}
	}
	if i != len(expected) {
		sr.Errorf("%s: found %d rows, expected %d", info, i, len(expected))
	}
}

func TestDebugName(t *testing.T) {
	defer leaktest.AfterTest(t)()
	s, db := setup(t)
	defer s.Stopper().Stop()

	file, _, _ := caller.Lookup(0)
	if err := db.Txn(context.TODO(), func(txn *client.Txn) error {
		if !strings.HasPrefix(txn.DebugName(), file+":") {
			t.Fatalf("expected \"%s\" to have the prefix \"%s:\"", txn.DebugName(), file)
		}
		return nil
	}); err != nil {
		t.Errorf("txn failed: %s", err)
	}
}

// logScope creates a testLogScope which corresponds to the
// lifetime of a logging directory. The logging directory is named
// after the caller of MaketestLogScope, up `skip` caller levels.
func logScope(t *testing.T) testLogScope {
	testName := "logUnknown"
	if _, _, f := caller.Lookup(1); f != "" {
		parts := strings.Split(f, ".")
		testName = "log" + parts[len(parts)-1]
	}
	tempDir, err := ioutil.TempDir("", testName)
	if err != nil {
		t.Fatal(err)
	}
	if err := dirTestOverride(tempDir); err != nil {
		t.Fatal(err)
	}
	return testLogScope(tempDir)
}

// RunAsyncTask runs function f in a goroutine. It returns an error when the
// Stopper is quiescing, in which case the function is not executed.
func (s *Stopper) RunAsyncTask(ctx context.Context, f func(context.Context)) error {
	file, line, _ := caller.Lookup(1)
	key := taskKey{file, line}
	if !s.runPrelude(key) {
		return errUnavailable
	}

	ctx, span := tracing.ForkCtxSpan(ctx, fmt.Sprintf("%s:%d", file, line))

	// Call f.
	go func() {
		defer s.Recover()
		defer s.runPostlude(key)
		defer tracing.FinishSpan(span)
		f(ctx)
	}()
	return nil
}

// rotateFile closes the syncBuffer's file and starts a new one.
func (sb *syncBuffer) rotateFile(now time.Time) error {
	if sb.file != nil {
		if err := sb.Flush(); err != nil {
			return err
		}
		if err := sb.file.Close(); err != nil {
			return err
		}
	}
	var err error
	sb.file, _, err = create(sb.sev, now)
	sb.nbytes = 0
	if err != nil {
		return err
	}

	sb.Writer = bufio.NewWriterSize(sb.file, bufferSize)

	f, l, _ := caller.Lookup(1)
	for _, msg := range []string{
		fmt.Sprintf("[config] file created at: %s\n", now.Format("2006/01/02 15:04:05")),
		fmt.Sprintf("[config] running on machine: %s\n", host),
		fmt.Sprintf("[config] binary: %s\n", build.GetInfo().Short()),
		fmt.Sprintf("[config] arguments: %s\n", os.Args),
		fmt.Sprintf("line format: [IWEF]yymmdd hh:mm:ss.uuuuuu goid file:line msg\n"),
	} {
		buf := formatLogEntry(Entry{
			Severity:  sb.sev,
			Time:      now.UnixNano(),
			Goroutine: goid.Get(),
			File:      f,
			Line:      int64(l),
			Message:   msg,
		}, nil, nil)
		var n int
		n, err = sb.file.Write(buf.Bytes())
		sb.nbytes += uint64(n)
		if err != nil {
			return err
		}
		logging.putBuffer(buf)
	}
	return nil
}

// StartCluster starts a cluster from the relevant flags. All test clusters
// should be created through this command since it sets up the logging in a
// unified way.
func StartCluster(t *testing.T, cfg cluster.TestConfig) (c cluster.Cluster) {
	var completed bool
	defer func() {
		if !completed && c != nil {
			c.AssertAndStop(t)
		}
	}()
	if !*flagRemote {
		logDir := *flagLogDir
		if logDir != "" {
			logDir = func(d string) string {
				for i := 1; i < 100; i++ {
					_, _, fun := caller.Lookup(i)
					if testFuncRE.MatchString(fun) {
						return filepath.Join(d, fun)
					}
				}
				panic("no caller matching Test(.*) in stack trace")
			}(logDir)
		}
		l := cluster.CreateLocal(cfg, logDir, *flagPrivileged, stopper)
		l.Start()
		c = l
		checkRangeReplication(t, l, 20*time.Second)
		completed = true
		return l
	}
	f := farmer(t, "")
	c = f
	if err := f.Resize(*flagNodes); err != nil {
		t.Fatal(err)
	}
	if err := f.WaitReady(5 * time.Minute); err != nil {
		if destroyErr := f.Destroy(t); destroyErr != nil {
			t.Fatalf("could not destroy cluster after error %v: %v", err, destroyErr)
		}
		t.Fatalf("cluster not ready in time: %v", err)
	}
	checkRangeReplication(t, f, 20*time.Second)
	completed = true
	return f
}

// MakeSrcCtx creates a SrcCtx value with contextual information about the
// caller at the requested depth.
func MakeSrcCtx(depth int) SrcCtx {
	f, l, fun := caller.Lookup(depth + 1)
	return SrcCtx{File: f, Line: l, Function: fun}
}

// StartCluster starts a cluster from the relevant flags. All test clusters
// should be created through this command since it sets up the logging in a
// unified way.
func StartCluster(t *testing.T, cfg cluster.TestConfig) (c cluster.Cluster) {
	var completed bool
	defer func() {
		if !completed && c != nil {
			c.AssertAndStop(t)
		}
	}()
	if *flagRemote {
		f := farmer(t, "")
		c = f
		if err := f.Resize(*flagNodes); err != nil {
			t.Fatal(err)
		}
		if err := f.WaitReady(5 * time.Minute); err != nil {
			if destroyErr := f.Destroy(t); destroyErr != nil {
				t.Fatalf("could not destroy cluster after error %s: %s", err, destroyErr)
			}
			t.Fatalf("cluster not ready in time: %s", err)
		}
	} else {
		logDir := *flagLogDir
		if logDir != "" {
			logDir = func(d string) string {
				for i := 1; i < 100; i++ {
					_, _, fun := caller.Lookup(i)
					if testFuncRE.MatchString(fun) {
						return filepath.Join(d, fun)
					}
				}
				panic("no caller matching Test(.*) in stack trace")
			}(logDir)
		}
		l := cluster.CreateLocal(cfg, logDir, *flagPrivileged, stopper)
		l.Start()
		c = l
	}
	wantedReplicas := 3
	if numNodes := c.NumNodes(); numNodes < wantedReplicas {
		wantedReplicas = numNodes
	}

	// Looks silly, but we actually start zero-node clusters in the
	// reference tests.
	if wantedReplicas > 0 {
		ctx := context.TODO()

		log.Infof(ctx, "waiting for first range to have %d replicas", wantedReplicas)

		util.SucceedsSoon(t, func() error {
			select {
			case <-stopper:
				t.Fatal("interrupted")
			case <-time.After(time.Second):
			}

			// Reconnect on every iteration; gRPC will eagerly tank the connection
			// on transport errors. Always talk to node 0 because it's guaranteed
			// to exist.
			client, dbStopper := c.NewClient(t, 0)
			defer dbStopper.Stop()

			ctxWithTimeout, cancel := context.WithTimeout(ctx, 5*time.Second)
			defer cancel()

			var desc roachpb.RangeDescriptor
			if err := client.GetProto(ctxWithTimeout, keys.RangeDescriptorKey(roachpb.RKeyMin), &desc); err != nil {
				return err
			}
			foundReplicas := len(desc.Replicas)

			if log.V(1) {
				log.Infof(ctxWithTimeout, "found %d replicas", foundReplicas)
			}

			if foundReplicas < wantedReplicas {
				return errors.Errorf("expected %d replicas, only found %d", wantedReplicas, foundReplicas)
			}
			return nil
		})
	}

	completed = true
	return c
}