// 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

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

	"golang.org/x/exp/trace/internal/event/go122"
)

// 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 {
	// State for coordinating with the recorder goroutine.
	fromTracer   *io.PipeReader
	toRecorder   *io.PipeWriter
	recorderWait sync.WaitGroup
	err          error

	// State specific to the recorder goroutine.
	header [16]byte
	active rawGeneration
	ringMu sync.Mutex
	ring   []rawGeneration

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

	enabled bool       // whether the flight recorder is enabled.
	writing sync.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
}

// 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
	r.fromTracer, r.toRecorder = io.Pipe()

	// Start tracing, sending data to the recorder goroutine (not yet started) via an io.Pipe.
	if err := trace.Start(r.toRecorder); err != nil {
		return err
	}

	// Start recorder goroutine.
	r.recorderWait.Add(1)
	go func() {
		defer r.recorderWait.Done()

		// Read in the header so we can tack it on to the front
		// of whatever WriteTo emits later.
		_, err := io.ReadFull(r.fromTracer, r.header[:])
		if err != nil {
			r.err = err
			return
		}

		// Process the rest of the trace.
		rd := bufio.NewReader(r.fromTracer)
		for {
			b, gen, err := readBatch(rd)
			if err == io.EOF || err == io.ErrClosedPipe {
				break
			}
			if err != nil {
				r.err = err
				return
			}

			// 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 {
					r.err = fmt.Errorf("broken trace: failed to find frequency event in generation %d", r.active.gen)
					return
				}

				// 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.isFreqBatch() {
				freq, err := parseFreq(b)
				if err != nil {
					r.err = err
					return
				}
				r.active.freq = freq
			}

			// Append the batch to the current generation.
			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)
		}
	}()

	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()

	// Close the write side of the pipe. This is safe because tracing has stopped, so no more will
	// be written to the pipe.
	r.fromTracer.Close()

	// Wait for the reader to exit.
	r.recorderWait.Wait()

	// Reset all state. No need to lock because the reader has already exited.
	r.active = rawGeneration{}
	r.ring = nil
	r.toRecorder.Close()
	r.fromTracer.Close()
	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.
	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 {
			// Rewrite the batch header event with four arguments: gen, M ID, timestamp, and data length.
			n, err := w.Write([]byte{byte(go122.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