// Copyright 2023 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.

//go:build go1.22

/*
Package trace provides a mechanism to collect and retrive
the most recent execution data without keeping the complete
execution tracing history.

The flight recorder was integrated into Go 1.25. The
integrated flight recorder should be used when possible.
*/
package trace

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"math/bits"
	"runtime/trace"
	"slices"
	stdsync "sync"
	"time"
	_ "unsafe" // for go:linkname

	"golang.org/x/exp/trace/internal/tracev2"
	"golang.org/x/exp/trace/internal/version"
)

// FlightRecorder represents a flight recording configuration.
//
// Flight recording holds execution trace data in a circular buffer representing
// the most recent execution data.
//
// Only one flight recording may be active at any given time.
type FlightRecorder struct {
	err error

	// State specific to the recorder.
	header  [16]byte
	active  rawGeneration
	ringMu  stdsync.Mutex
	ring    []rawGeneration
	freq    frequency // timestamp conversion factor, from the runtime
	version version.Version

	// Externally-set options.
	targetSize   int
	targetPeriod time.Duration

	enabled bool          // whether the flight recorder is enabled.
	writing stdsync.Mutex // protects concurrent calls to WriteTo

	// The values of targetSize and targetPeriod we've committed to since the last Start.
	wantSize int
	wantDur  time.Duration
}

// NewFlightRecorder creates a new flight recording configuration.
func NewFlightRecorder() *FlightRecorder {
	return &FlightRecorder{
		// These are just some optimistic, reasonable defaults.
		//
		// In reality we're also bound by whatever the runtime defaults are, because
		// we currently have no way to change them.
		//
		// TODO(mknyszek): Consider adding a function that allows mutating one or
		// both of these values' equivalents in the runtime.
		targetSize:   10 << 20, // 10 MiB.
		targetPeriod: 10 * time.Second,
	}
}

// SetPeriod sets the approximate time duration that the flight recorder's circular buffer
// represents.
//
// Note that SetPeriod does not make any guarantees on the amount of time the trace
// produced by WriteTo will represent.
// This is just a hint to the runtime to enable some control the resulting trace.
//
// The initial period is implementation defined, but can be assumed to be on the order
// of seconds.
//
// Adjustments to this value will not apply to an active flight recorder, and will not apply
// if tracing is already enabled via trace.Start. All tracing must be stopped and started
// again to change this value.
func (r *FlightRecorder) SetPeriod(d time.Duration) {
	r.targetPeriod = d
}

// SetSize sets the approximate size of the flight recorder's circular buffer.
//
// This generally takes precedence over the duration passed to SetPeriod.
// However, it does not make any guarantees on the size of the data WriteTo will write.
// This is just a hint to the runtime to enable some control over the memory overheads
// of tracing.
//
// The initial size is implementation defined.
//
// Adjustments to this value will not apply to an active flight recorder, and will not apply
// if tracing is already enabled via trace.Start. All tracing must be stopped and started
// again to change this value.
func (r *FlightRecorder) SetSize(bytes int) {
	r.targetSize = bytes
}

// A recorder receives bytes from the runtime tracer, processes it.
type recorder struct {
	r *FlightRecorder

	headerReceived bool
}

func (w *recorder) Write(p []byte) (n int, err error) {
	r := w.r

	defer func() {
		if err != nil {
			// Propagate errors to the flightrecorder.
			if r.err == nil {
				r.err = err
			}
			trace.Stop() // Stop the tracer, preventing further writes.
		}
	}()

	rd := bytes.NewReader(p)

	if !w.headerReceived {
		if len(p) < len(r.header) {
			return 0, fmt.Errorf("expected at least %d bytes in the first write", len(r.header))
		}
		rd.Read(r.header[:])
		if r.version, err = version.ReadHeader(bytes.NewReader(r.header[:])); err != nil {
			return len(p) - rd.Len(), fmt.Errorf("failed to parse version from header: %s", err)
		}
		w.headerReceived = true
	}

	b, gen, err := readBatch(rd) // Every write from the runtime is guaranteed to be a complete batch.
	if err == io.EOF {
		if rd.Len() > 0 {
			return len(p) - rd.Len(), errors.New("short read")
		}
		return len(p), nil
	}
	if err != nil {
		return len(p) - rd.Len(), err
	}
	if b.isEndOfGeneration() {
		gen = r.active.gen
	}

	// Check if we're entering a new generation.
	if r.active.gen != 0 && r.active.gen+1 == gen {
		r.ringMu.Lock()

		// Validate r.active.freq before we use it. It's required for a generation
		// to not be considered broken, and without it, we can't correctly handle
		// SetPeriod.
		if r.active.freq == 0 {
			return len(p) - rd.Len(), fmt.Errorf("broken trace: failed to find frequency event in generation %d", r.active.gen)
		}

		// Get the current trace clock time.
		now := traceTimeNow(r.active.freq)

		// Add the current generation to the ring. Make sure we always have at least one
		// complete generation by putting the active generation onto the new list, regardless
		// of whatever our settings are.
		//
		// N.B. Let's completely replace the ring here, so that WriteTo can just make a copy
		// and not worry about aliasing. This creates allocations, but at a very low rate.
		newRing := []rawGeneration{r.active}
		size := r.active.size
		for i := len(r.ring) - 1; i >= 0; i-- {
			// Stop adding older generations if the new ring already exceeds the thresholds.
			// This ensures we keep generations that cross a threshold, but not any that lie
			// entirely outside it.
			if size > r.wantSize || now.Sub(newRing[len(newRing)-1].minTraceTime()) > r.wantDur {
				break
			}
			size += r.ring[i].size
			newRing = append(newRing, r.ring[i])
		}
		slices.Reverse(newRing)
		r.ring = newRing
		r.ringMu.Unlock()

		// Start a new active generation.
		r.active = rawGeneration{}
	}

	// Obtain the frequency if this is a frequency batch.
	if b.isSyncBatch(r.version) {
		var s sync
		setSyncBatch(&s, b, r.version)
		r.active.freq = s.freq
	}

	// Append the batch to the current generation.
	if !b.isEndOfGeneration() {
		if r.active.gen == 0 {
			r.active.gen = gen
		}
		if r.active.minTime == 0 || r.active.minTime > b.time {
			r.active.minTime = b.time
		}
		r.active.size += 1
		r.active.size += uvarintSize(gen)
		r.active.size += uvarintSize(uint64(b.m))
		r.active.size += uvarintSize(uint64(b.time))
		r.active.size += uvarintSize(uint64(len(b.data)))
		r.active.size += len(b.data)
	}
	r.active.batches = append(r.active.batches, b)

	return len(p) - rd.Len(), nil
}

// Start begins flight recording. Only one flight recorder or one call to [runtime/trace.Start]
// may be active at any given time. Returns an error if starting the flight recorder would
// violate this rule.
func (r *FlightRecorder) Start() error {
	if r.enabled {
		return fmt.Errorf("cannot enable a enabled flight recorder")
	}

	r.wantSize = r.targetSize
	r.wantDur = r.targetPeriod
	r.err = nil

	// Start tracing, data is sent to a recorder which forwards it to our own
	// storage.
	if err := trace.Start(&recorder{r: r}); err != nil {
		return err
	}

	r.enabled = true
	return nil
}

// Stop ends flight recording. It waits until any concurrent [FlightRecorder.WriteTo] calls exit.
// Returns an error if the flight recorder is inactive.
func (r *FlightRecorder) Stop() error {
	if !r.enabled {
		return fmt.Errorf("cannot disable a disabled flight recorder")
	}
	r.enabled = false
	trace.Stop()

	// Reset all state. No need to lock because the reader has already exited.
	r.active = rawGeneration{}
	r.ring = nil
	return r.err
}

// Enabled returns true if the flight recorder is active. Specifically, it will return true if
// Start did not return an error, and Stop has not yet been called.
// It is safe to call from multiple goroutines simultaneously.
func (r *FlightRecorder) Enabled() bool {
	return r.enabled
}

// ErrSnapshotActive indicates that a call to WriteTo was made while one was already in progress.
// If the caller of WriteTo sees this error, they should use the result from the other call to WriteTo.
var ErrSnapshotActive = fmt.Errorf("call to WriteTo for trace.FlightRecorder already in progress")

// WriteTo takes a snapshots of the circular buffer's contents and writes the execution data to w.
// Returns the number of bytes written and an error.
// An error is returned upon failure to write to w or if the flight recorder is inactive.
// Only one goroutine may execute WriteTo at a time, but it is safe to call from multiple goroutines.
// If a goroutine calls WriteTo while another goroutine is currently executing it, WriteTo will return
// ErrSnapshotActive to that goroutine.
func (r *FlightRecorder) WriteTo(w io.Writer) (total int, err error) {
	if !r.enabled {
		return 0, fmt.Errorf("cannot snapshot a disabled flight recorder")
	}
	if !r.writing.TryLock() {
		return 0, ErrSnapshotActive
	}
	defer r.writing.Unlock()

	// Force a global buffer flush twice.
	//
	// This is pretty unfortunate, but because the signal that a generation is done is that a new
	// generation appears in the trace *or* the trace stream ends, the recorder goroutine will
	// have no idea when to add a generation to the ring if we just flush once. If we flush twice,
	// at least the first one will end up on the ring, which is the one we wanted anyway.
	//
	// In a runtime-internal implementation this is a non-issue. The runtime is fully aware
	// of what generations are complete, so only one flush is necessary.
	//
	// As of Go 1.25, it's not required to call runtime_traceAdvance twice.
	runtime_traceAdvance(false)
	runtime_traceAdvance(false)

	// Now that everything has been flushed and written, grab whatever we have.
	//
	// N.B. traceAdvance blocks until the tracer goroutine has actually written everything
	// out, which means the generation we just flushed must have been already been observed
	// by the recorder goroutine. Because we flushed twice, the first flush is guaranteed to
	// have been both completed *and* processed by the recorder goroutine.
	r.ringMu.Lock()
	gens := r.ring
	r.ringMu.Unlock()

	// Write the header.
	total, err = w.Write(r.header[:])
	if err != nil {
		return total, err
	}

	// Helper for writing varints.
	var varintBuf [binary.MaxVarintLen64]byte
	writeUvarint := func(u uint64) error {
		v := binary.PutUvarint(varintBuf[:], u)
		n, err := w.Write(varintBuf[:v])
		total += n
		return err
	}

	// Write all the data.
	for _, gen := range gens {
		for _, batch := range gen.batches {
			var n int
			if !batch.isEndOfGeneration() {
				// Rewrite the batch header event with four arguments: gen, M ID, timestamp, and data length.
				n, err := w.Write([]byte{byte(tracev2.EvEventBatch)})
				total += n
				if err != nil {
					return total, err
				}
				if err := writeUvarint(gen.gen); err != nil {
					return total, err
				}
				if err := writeUvarint(uint64(batch.m)); err != nil {
					return total, err
				}
				if err := writeUvarint(uint64(batch.time)); err != nil {
					return total, err
				}
				if err := writeUvarint(uint64(len(batch.data))); err != nil {
					return total, err
				}
			}

			// Write batch data.
			n, err = w.Write(batch.data)
			total += n
			if err != nil {
				return total, err
			}
		}
	}
	return total, nil
}

type rawGeneration struct {
	gen     uint64
	size    int
	minTime timestamp
	freq    frequency
	batches []batch
}

func (r *rawGeneration) minTraceTime() Time {
	return r.freq.mul(r.minTime)
}

func traceTimeNow(freq frequency) Time {
	// TODO(mknyszek): It's unfortunate that we have to rely on runtime-internal details
	// like this. This would be better off in the runtime.
	return freq.mul(timestamp(runtime_traceClockNow()))
}

func uvarintSize(x uint64) int {
	return 1 + bits.Len64(x)/7
}

//go:linkname runtime_traceAdvance runtime.traceAdvance
func runtime_traceAdvance(stopTrace bool)

//go:linkname runtime_traceClockNow runtime.traceClockNow
func runtime_traceClockNow() int64