checkpoint.go

// Copyright 2021 PingCAP, 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package checkpoint

import (
	"bufio"
	"context"
	"encoding/json"
	"io"
	"os"
	"runtime"
	"strings"
	"sync/atomic"

	"github.com/pingcap/errors"
	"github.com/pingcap/tiup/pkg/logger/log"
	"github.com/pingcap/tiup/pkg/version"
	"go.uber.org/zap"
	"go.uber.org/zap/zapcore"
	"golang.org/x/sync/semaphore"
)

type contextKey string

const (
	semKey       = contextKey("CHECKPOINT_SEMAPHORE")
	goroutineKey = contextKey("CHECKPOINT_GOROUTINE")
	funcKey      = "__func__"
	hashKey      = "__hash__"

	// At most 10M for each line in audit log
	maxTokenSize = 10 * 1024 * 1024
)

var (
	checkpoint *CheckPoint

	// DebugCheckpoint is a switch used to debug if:
	// - The context passed to checkpoint is generated by checkpoint.NewContext
	// - multilple context acquire applied to the same context belone to the sampe goroutine
	DebugCheckpoint = os.Getenv("DEBUG_CHECKPOINT") == "1"
)

// SetCheckPoint set global checkpoint for replay
func SetCheckPoint(file string) error {
	pointReader, err := os.Open(file)
	if err != nil {
		return errors.AddStack(err)
	}
	defer pointReader.Close()

	checkpoint, err = NewCheckPoint(pointReader)
	if err != nil {
		return err
	}

	return nil
}

// HasCheckPoint returns if SetCheckPoint has been called
func HasCheckPoint() bool {
	return checkpoint != nil
}

// Acquire wraps CheckPoint.Acquire
func Acquire(ctx context.Context, fs FieldSet, point map[string]interface{}) *Point {
	if ctx.Value(goroutineKey) == nil || ctx.Value(semKey) == nil {
		if DebugCheckpoint {
			panic("the context passed to checkpoint.Acquire is not generated by checkpoint.NewContext")
		}
		log.Debugf("context missing for checkpoint, the result of replaying this operation may be unexpected!")
		ctx = NewContext(ctx)
	}

	// Check goroutine if we are in test
	gptr := ctx.Value(goroutineKey).(*goroutineLock)
	g := atomic.LoadUint64((*uint64)(gptr))
	if g == 0 {
		atomic.StoreUint64((*uint64)(gptr), uint64(newGoroutineLock()))
	} else {
		goroutineLock(g).check()
	}

	pc, _, _, _ := runtime.Caller(1)
	fn := runtime.FuncForPC(pc).Name()

	// If checkpoint is disabled, return a mock point
	if checkpoint == nil {
		return &Point{nil, fn, nil, true}
	}

	return checkpoint.acquire(ctx, fs, fn, point)
}

// NewContext wraps given context with value needed by checkpoint
func NewContext(ctx context.Context) context.Context {
	switch {
	case ctx.Value(semKey) == nil:
		ctx = context.WithValue(ctx, semKey, semaphore.NewWeighted(1))
	case ctx.Value(semKey).(*semaphore.Weighted).TryAcquire(1):
		defer ctx.Value(semKey).(*semaphore.Weighted).Release(1)
		ctx = context.WithValue(ctx, semKey, semaphore.NewWeighted(1))
	default:
		ctx = context.WithValue(ctx, semKey, semaphore.NewWeighted(0))
	}

	return context.WithValue(ctx, goroutineKey, new(goroutineLock))
}

// CheckPoint provides the ability to recover from a failed command at the failpoint
type CheckPoint struct {
	points []map[string]interface{}
}

// NewCheckPoint returns a CheckPoint by given audit file
func NewCheckPoint(r io.Reader) (*CheckPoint, error) {
	cp := CheckPoint{points: make([]map[string]interface{}, 0)}

	scanner := bufio.NewScanner(r)
	scanner.Buffer(nil, maxTokenSize)
	for scanner.Scan() {
		line := scanner.Text()
		m, err := checkLine(line)
		if err != nil {
			return nil, errors.Annotate(err, "initial checkpoint failed")
		}
		if m == nil {
			continue
		}

		cp.points = append(cp.points, m)
	}

	if err := scanner.Err(); err != nil {
		return nil, errors.Annotate(err, "failed to parse audit file %s")
	}

	return &cp, nil
}

// Acquire get point from checkpoints
func (c *CheckPoint) acquire(ctx context.Context, fs FieldSet, fn string, point map[string]interface{}) *Point {
	acquired := ctx.Value(semKey).(*semaphore.Weighted).TryAcquire(1)
	point[funcKey] = fn
	point[hashKey] = version.GitHash

next_point:
	for _, p := range c.points {
		for _, cf := range fs.Slice() {
			if cf.eq == nil {
				continue
			}
			if !contains(p, cf.field) || !contains(point, cf.field) || !cf.eq(p[cf.field], point[cf.field]) {
				continue next_point
			}
		}
		return &Point{ctx, fn, p, acquired}
	}
	return &Point{ctx, fn, nil, acquired}
}

// Point is a point of checkpoint
type Point struct {
	ctx      context.Context
	fn       string
	point    map[string]interface{}
	acquired bool
}

// Hit returns value of the point, it will be nil if not hit.
func (p *Point) Hit() map[string]interface{} {
	return p.point
}

// Release write checkpoint into log file
func (p *Point) Release(err error, fields ...zapcore.Field) {
	logfn := zap.L().Info
	if err != nil {
		logfn = zap.L().Error
		fields = append(fields, zap.Error(err))
	}
	fields = append(fields,
		zap.String(hashKey, version.GitHash),
		zap.String(funcKey, p.fn),
		zap.Bool("hit", p.Hit() != nil))

	if p.acquired {
		logfn("CheckPoint", fields...)
		// If checkpoint is disabled, the p.ctx will be nil
		if p.ctx != nil {
			p.ctx.Value(semKey).(*semaphore.Weighted).Release(1)
		}
	}
}

func checkLine(line string) (map[string]interface{}, error) {
	// target log format:
	//	2021-01-13T14:11:02.987+0800    INFO    SCPCommand      {k:v...}
	//	2021-01-13T14:11:03.780+0800    INFO    SSHCommand      {k:v...}
	ss := strings.Fields(line)
	pos := strings.Index(line, "{")
	if len(ss) < 4 || ss[1] != "INFO" || ss[2] != "CheckPoint" || pos == -1 {
		return nil, nil
	}

	m := make(map[string]interface{})
	if err := json.Unmarshal([]byte(line[pos:]), &m); err != nil {
		return nil, errors.AddStack(err)
	}

	return m, nil
}

func contains(m map[string]interface{}, f string) bool {
	_, ok := m[f]
	return ok
}