Implement trace archiving and management features
- Added functionality to archive traces by compressing and marking them as read-only, enhancing data management capabilities. - Introduced methods for checking if a trace is archived and for unarchiving traces, allowing for better control over trace data. - Updated the trace manager and driver interfaces to support archiving operations, ensuring consistency across components. - Enhanced error handling for operations on archived traces, preventing modifications to read-only data. - Improved trace metadata structure to include archived status and timestamp, facilitating better trace lifecycle management.
This commit is contained in:
Max
parent
396a37d32f
commit
870691e3d4
8 changed files with 739 additions and 25 deletions
|
|
@ -177,7 +177,7 @@ func (ctx *Context) Trace() (traceTypes.Manager, error) {
|
|||
}
|
||||
|
||||
// Prepare trace options
|
||||
traceOption := &traceTypes.TraceOption{ID: traceID}
|
||||
traceOption := &traceTypes.TraceOption{ID: traceID, AutoArchive: config.Conf.Mode == "production"}
|
||||
|
||||
// Set trace options from authorized information
|
||||
if ctx.Authorized != nil {
|
||||
|
|
|
|||
|
|
@ -1,12 +1,16 @@
|
|||
package local
|
||||
|
||||
import (
|
||||
"archive/tar"
|
||||
"compress/gzip"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/yaoapp/yao/config"
|
||||
"github.com/yaoapp/yao/trace/types"
|
||||
|
|
@ -133,6 +137,15 @@ func (d *Driver) ensureTraceDir(traceID string) error {
|
|||
|
||||
// SaveNode persists a node to disk
|
||||
func (d *Driver) SaveNode(ctx context.Context, traceID string, node *types.TraceNode) error {
|
||||
// Check if archived - archived traces are read-only
|
||||
archived, err := d.IsArchived(ctx, traceID)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to check archive status: %w", err)
|
||||
}
|
||||
if archived {
|
||||
return fmt.Errorf("cannot save node: trace %s is archived (read-only)", traceID)
|
||||
}
|
||||
|
||||
if err := d.ensureTraceDir(traceID); err != nil {
|
||||
return err
|
||||
}
|
||||
|
|
@ -162,6 +175,18 @@ func (d *Driver) SaveNode(ctx context.Context, traceID string, node *types.Trace
|
|||
|
||||
// LoadNode loads a node from disk
|
||||
func (d *Driver) LoadNode(ctx context.Context, traceID string, nodeID string) (*types.TraceNode, error) {
|
||||
// Check if archived and extract if needed
|
||||
archived, err := d.IsArchived(ctx, traceID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to check archive status: %w", err)
|
||||
}
|
||||
if archived {
|
||||
// Extract archive for read access
|
||||
if err := d.unarchive(ctx, traceID); err != nil {
|
||||
return nil, fmt.Errorf("failed to unarchive trace: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
filePath := filepath.Join(d.getTracePath(traceID), "nodes", nodeID+".json")
|
||||
|
||||
data, err := os.ReadFile(filePath)
|
||||
|
|
@ -215,6 +240,15 @@ func (d *Driver) LoadTrace(ctx context.Context, traceID string) (*types.TraceNod
|
|||
|
||||
// SaveSpace persists a space to disk
|
||||
func (d *Driver) SaveSpace(ctx context.Context, traceID string, space *types.TraceSpace) error {
|
||||
// Check if archived - archived traces are read-only
|
||||
archived, err := d.IsArchived(ctx, traceID)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to check archive status: %w", err)
|
||||
}
|
||||
if archived {
|
||||
return fmt.Errorf("cannot save space: trace %s is archived (read-only)", traceID)
|
||||
}
|
||||
|
||||
if err := d.ensureTraceDir(traceID); err != nil {
|
||||
return err
|
||||
}
|
||||
|
|
@ -241,6 +275,17 @@ func (d *Driver) SaveSpace(ctx context.Context, traceID string, space *types.Tra
|
|||
|
||||
// LoadSpace loads a space from disk
|
||||
func (d *Driver) LoadSpace(ctx context.Context, traceID string, spaceID string) (*types.TraceSpace, error) {
|
||||
// Check if archived and extract if needed
|
||||
archived, err := d.IsArchived(ctx, traceID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to check archive status: %w", err)
|
||||
}
|
||||
if archived {
|
||||
if err := d.unarchive(ctx, traceID); err != nil {
|
||||
return nil, fmt.Errorf("failed to unarchive trace: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
filePath := filepath.Join(d.getTracePath(traceID), "spaces", spaceID+".json")
|
||||
|
||||
data, err := os.ReadFile(filePath)
|
||||
|
|
@ -423,6 +468,15 @@ func (d *Driver) ListSpaceKeys(ctx context.Context, traceID, spaceID string) ([]
|
|||
|
||||
// SaveLog appends a log entry to disk
|
||||
func (d *Driver) SaveLog(ctx context.Context, traceID string, log *types.TraceLog) error {
|
||||
// Check if archived - archived traces are read-only
|
||||
archived, err := d.IsArchived(ctx, traceID)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to check archive status: %w", err)
|
||||
}
|
||||
if archived {
|
||||
return fmt.Errorf("cannot save log: trace %s is archived (read-only)", traceID)
|
||||
}
|
||||
|
||||
if err := d.ensureTraceDir(traceID); err != nil {
|
||||
return err
|
||||
}
|
||||
|
|
@ -458,6 +512,17 @@ func (d *Driver) SaveLog(ctx context.Context, traceID string, log *types.TraceLo
|
|||
|
||||
// LoadLogs loads all logs for a trace or specific node from disk
|
||||
func (d *Driver) LoadLogs(ctx context.Context, traceID string, nodeID string) ([]*types.TraceLog, error) {
|
||||
// Check if archived and extract if needed
|
||||
archived, err := d.IsArchived(ctx, traceID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to check archive status: %w", err)
|
||||
}
|
||||
if archived {
|
||||
if err := d.unarchive(ctx, traceID); err != nil {
|
||||
return nil, fmt.Errorf("failed to unarchive trace: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
logsDir := filepath.Join(d.getTracePath(traceID), "logs")
|
||||
|
||||
var logs []*types.TraceLog
|
||||
|
|
@ -527,6 +592,16 @@ func (d *Driver) loadLogFile(filePath string) ([]*types.TraceLog, error) {
|
|||
|
||||
// SaveTraceInfo persists trace metadata to disk
|
||||
func (d *Driver) SaveTraceInfo(ctx context.Context, info *types.TraceInfo) error {
|
||||
// Allow saving trace info even if archived (for updating archive status)
|
||||
// But check if it's trying to modify a non-archive field
|
||||
if info.Archived {
|
||||
// If already archived, only allow updating archive-related fields
|
||||
existing, err := d.LoadTraceInfo(ctx, info.ID)
|
||||
if err == nil && existing != nil && existing.Archived && !info.Archived {
|
||||
return fmt.Errorf("cannot unarchive trace: trace %s is archived", info.ID)
|
||||
}
|
||||
}
|
||||
|
||||
if err := d.ensureTraceDir(info.ID); err != nil {
|
||||
return err
|
||||
}
|
||||
|
|
@ -545,8 +620,8 @@ func (d *Driver) SaveTraceInfo(ctx context.Context, info *types.TraceInfo) error
|
|||
return nil
|
||||
}
|
||||
|
||||
// LoadTraceInfo loads trace metadata from disk
|
||||
func (d *Driver) LoadTraceInfo(ctx context.Context, traceID string) (*types.TraceInfo, error) {
|
||||
// loadTraceInfoDirect loads trace info without unarchiving (internal use)
|
||||
func (d *Driver) loadTraceInfoDirect(traceID string) (*types.TraceInfo, error) {
|
||||
filePath := filepath.Join(d.getTracePath(traceID), "trace_info.json")
|
||||
|
||||
data, err := os.ReadFile(filePath)
|
||||
|
|
@ -565,6 +640,22 @@ func (d *Driver) LoadTraceInfo(ctx context.Context, traceID string) (*types.Trac
|
|||
return &info, nil
|
||||
}
|
||||
|
||||
// LoadTraceInfo loads trace metadata from disk
|
||||
func (d *Driver) LoadTraceInfo(ctx context.Context, traceID string) (*types.TraceInfo, error) {
|
||||
// Check if archived and extract if needed
|
||||
archived, err := d.IsArchived(ctx, traceID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to check archive status: %w", err)
|
||||
}
|
||||
if archived {
|
||||
if err := d.unarchive(ctx, traceID); err != nil {
|
||||
return nil, fmt.Errorf("failed to unarchive trace: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
return d.loadTraceInfoDirect(traceID)
|
||||
}
|
||||
|
||||
// DeleteTrace removes entire trace from disk
|
||||
func (d *Driver) DeleteTrace(ctx context.Context, traceID string) error {
|
||||
tracePath := d.getTracePath(traceID)
|
||||
|
|
@ -578,6 +669,15 @@ func (d *Driver) DeleteTrace(ctx context.Context, traceID string) error {
|
|||
|
||||
// SaveUpdate persists a trace update event to disk (append-only)
|
||||
func (d *Driver) SaveUpdate(ctx context.Context, traceID string, update *types.TraceUpdate) error {
|
||||
// Check if archived - archived traces are read-only
|
||||
archived, err := d.IsArchived(ctx, traceID)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to check archive status: %w", err)
|
||||
}
|
||||
if archived {
|
||||
return fmt.Errorf("cannot save update: trace %s is archived (read-only)", traceID)
|
||||
}
|
||||
|
||||
if err := d.ensureTraceDir(traceID); err != nil {
|
||||
return err
|
||||
}
|
||||
|
|
@ -641,6 +741,205 @@ func (d *Driver) LoadUpdates(ctx context.Context, traceID string, since int64) (
|
|||
return updates, nil
|
||||
}
|
||||
|
||||
// Archive archives a trace by compressing it to tar.gz
|
||||
func (d *Driver) Archive(ctx context.Context, traceID string) error {
|
||||
tracePath := d.getTracePath(traceID)
|
||||
archivePath := tracePath + ".tar.gz"
|
||||
archivedMarker := filepath.Join(filepath.Dir(tracePath), "."+traceID+".archived")
|
||||
|
||||
// Check if already archived
|
||||
if _, err := os.Stat(archivedMarker); err == nil {
|
||||
return fmt.Errorf("trace %s is already archived", traceID)
|
||||
}
|
||||
|
||||
// Check if trace directory exists
|
||||
if _, err := os.Stat(tracePath); os.IsNotExist(err) {
|
||||
return fmt.Errorf("trace directory not found: %s", traceID)
|
||||
}
|
||||
|
||||
// Update trace info to mark as archived BEFORE creating archive
|
||||
info, err := d.loadTraceInfoDirect(traceID)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to load trace info: %w", err)
|
||||
}
|
||||
if info != nil {
|
||||
now := time.Now().UnixMilli()
|
||||
info.Archived = true
|
||||
info.ArchivedAt = &now
|
||||
// Write trace info back before archiving
|
||||
infoPath := filepath.Join(tracePath, "trace_info.json")
|
||||
infoData, err := json.MarshalIndent(info, "", " ")
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to marshal trace info: %w", err)
|
||||
}
|
||||
if err := os.WriteFile(infoPath, infoData, 0644); err != nil {
|
||||
return fmt.Errorf("failed to write trace info: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
// Create tar.gz archive
|
||||
archiveFile, err := os.Create(archivePath)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to create archive file: %w", err)
|
||||
}
|
||||
defer archiveFile.Close()
|
||||
|
||||
// Create gzip writer
|
||||
gzipWriter := gzip.NewWriter(archiveFile)
|
||||
defer gzipWriter.Close()
|
||||
|
||||
// Create tar writer
|
||||
tarWriter := tar.NewWriter(gzipWriter)
|
||||
defer tarWriter.Close()
|
||||
|
||||
// Walk through trace directory and add files to archive
|
||||
err = filepath.Walk(tracePath, func(file string, fi os.FileInfo, err error) error {
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Create tar header
|
||||
header, err := tar.FileInfoHeader(fi, fi.Name())
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Update header name to be relative to trace directory
|
||||
relPath, err := filepath.Rel(tracePath, file)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
header.Name = relPath
|
||||
|
||||
// Write header
|
||||
if err := tarWriter.WriteHeader(header); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// If not a directory, write file content
|
||||
if !fi.IsDir() {
|
||||
data, err := os.Open(file)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer data.Close()
|
||||
|
||||
if _, err := io.Copy(tarWriter, data); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
})
|
||||
|
||||
if err != nil {
|
||||
// Clean up failed archive
|
||||
os.Remove(archivePath)
|
||||
return fmt.Errorf("failed to create archive: %w", err)
|
||||
}
|
||||
|
||||
// Create archived marker file
|
||||
if err := os.WriteFile(archivedMarker, []byte(time.Now().Format(time.RFC3339)), 0644); err != nil {
|
||||
return fmt.Errorf("failed to create archived marker: %w", err)
|
||||
}
|
||||
|
||||
// Remove original directory
|
||||
if err := os.RemoveAll(tracePath); err != nil {
|
||||
return fmt.Errorf("failed to remove original trace directory: %w", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// IsArchived checks if a trace is archived
|
||||
func (d *Driver) IsArchived(ctx context.Context, traceID string) (bool, error) {
|
||||
tracePath := d.getTracePath(traceID)
|
||||
archivedMarker := filepath.Join(filepath.Dir(tracePath), "."+traceID+".archived")
|
||||
|
||||
// Check marker file
|
||||
if _, err := os.Stat(archivedMarker); err == nil {
|
||||
return true, nil
|
||||
}
|
||||
|
||||
// Also check if archive file exists
|
||||
archivePath := tracePath + ".tar.gz"
|
||||
if _, err := os.Stat(archivePath); err == nil {
|
||||
return true, nil
|
||||
}
|
||||
|
||||
return false, nil
|
||||
}
|
||||
|
||||
// unarchive extracts an archived trace (helper method, not exposed in interface)
|
||||
func (d *Driver) unarchive(ctx context.Context, traceID string) error {
|
||||
tracePath := d.getTracePath(traceID)
|
||||
archivePath := tracePath + ".tar.gz"
|
||||
|
||||
// Check if archive exists
|
||||
if _, err := os.Stat(archivePath); os.IsNotExist(err) {
|
||||
return fmt.Errorf("archive not found: %s", traceID)
|
||||
}
|
||||
|
||||
// Open archive file
|
||||
archiveFile, err := os.Open(archivePath)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to open archive: %w", err)
|
||||
}
|
||||
defer archiveFile.Close()
|
||||
|
||||
// Create gzip reader
|
||||
gzipReader, err := gzip.NewReader(archiveFile)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to create gzip reader: %w", err)
|
||||
}
|
||||
defer gzipReader.Close()
|
||||
|
||||
// Create tar reader
|
||||
tarReader := tar.NewReader(gzipReader)
|
||||
|
||||
// Extract files
|
||||
for {
|
||||
header, err := tarReader.Next()
|
||||
if err == io.EOF {
|
||||
break
|
||||
}
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to read tar header: %w", err)
|
||||
}
|
||||
|
||||
// Construct target path
|
||||
target := filepath.Join(tracePath, header.Name)
|
||||
|
||||
// Create directory if needed
|
||||
if header.Typeflag == tar.TypeDir {
|
||||
if err := os.MkdirAll(target, 0755); err != nil {
|
||||
return fmt.Errorf("failed to create directory: %w", err)
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
// Create parent directory
|
||||
if err := os.MkdirAll(filepath.Dir(target), 0755); err != nil {
|
||||
return fmt.Errorf("failed to create parent directory: %w", err)
|
||||
}
|
||||
|
||||
// Create file
|
||||
outFile, err := os.Create(target)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to create file: %w", err)
|
||||
}
|
||||
|
||||
// Copy file content
|
||||
if _, err := io.Copy(outFile, tarReader); err != nil {
|
||||
outFile.Close()
|
||||
return fmt.Errorf("failed to write file: %w", err)
|
||||
}
|
||||
outFile.Close()
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// Close closes the local driver
|
||||
func (d *Driver) Close() error {
|
||||
// No cleanup needed for local file system
|
||||
|
|
|
|||
|
|
@ -16,19 +16,27 @@ type manager struct {
|
|||
traceID string
|
||||
driver types.Driver
|
||||
stateCmdChan chan stateCommand // Single channel for all state mutations
|
||||
autoArchive bool // Auto-archive on complete/fail
|
||||
}
|
||||
|
||||
// NewManager creates a new trace manager instance
|
||||
func NewManager(ctx context.Context, traceID string, driver types.Driver) (types.Manager, error) {
|
||||
func NewManager(ctx context.Context, traceID string, driver types.Driver, option *types.TraceOption) (types.Manager, error) {
|
||||
// Create a cancellable context for the manager
|
||||
managerCtx, cancel := context.WithCancel(ctx)
|
||||
|
||||
// Determine auto-archive setting
|
||||
autoArchive := false
|
||||
if option != nil {
|
||||
autoArchive = option.AutoArchive
|
||||
}
|
||||
|
||||
m := &manager{
|
||||
ctx: managerCtx,
|
||||
cancel: cancel,
|
||||
traceID: traceID,
|
||||
driver: driver,
|
||||
stateCmdChan: make(chan stateCommand, 100), // Buffered channel for performance
|
||||
autoArchive: autoArchive,
|
||||
}
|
||||
|
||||
// Start state worker goroutine
|
||||
|
|
@ -595,6 +603,17 @@ func (m *manager) MarkComplete() error {
|
|||
Data: types.NewTraceCompleteData(m.traceID, totalDuration),
|
||||
})
|
||||
|
||||
// Auto-archive if enabled
|
||||
if m.autoArchive {
|
||||
if err := m.driver.Archive(m.ctx, m.traceID); err != nil {
|
||||
// Log error but don't fail the complete operation
|
||||
m.Debug("Failed to auto-archive trace", map[string]any{
|
||||
"trace_id": m.traceID,
|
||||
"error": err.Error(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -1,11 +1,15 @@
|
|||
package store
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"compress/gzip"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/yaoapp/gou/store"
|
||||
"github.com/yaoapp/yao/trace/types"
|
||||
|
|
@ -127,8 +131,32 @@ func (d *Driver) getKey(traceID string, parts ...string) string {
|
|||
return strings.Join(allParts, ":")
|
||||
}
|
||||
|
||||
// getKeyPrefix returns the prefix for all keys of a trace
|
||||
func (d *Driver) getKeyPrefix(traceID string) string {
|
||||
return d.prefix + ":" + traceID + ":"
|
||||
}
|
||||
|
||||
// getTraceInfoKey returns the key for trace info
|
||||
func (d *Driver) getTraceInfoKey(traceID string) string {
|
||||
return d.getKey(traceID, "info")
|
||||
}
|
||||
|
||||
// getUpdatesKey returns the key for trace updates
|
||||
func (d *Driver) getUpdatesKey(traceID string) string {
|
||||
return d.getKey(traceID, "updates")
|
||||
}
|
||||
|
||||
// SaveNode persists a node to store
|
||||
func (d *Driver) SaveNode(ctx context.Context, traceID string, node *types.TraceNode) error {
|
||||
// Check if archived - archived traces are read-only
|
||||
archived, err := d.IsArchived(ctx, traceID)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to check archive status: %w", err)
|
||||
}
|
||||
if archived {
|
||||
return fmt.Errorf("cannot save node: trace %s is archived (read-only)", traceID)
|
||||
}
|
||||
|
||||
key := d.getKey(traceID, "node", node.ID)
|
||||
|
||||
// Convert to persist format (only store children IDs)
|
||||
|
|
@ -148,6 +176,17 @@ func (d *Driver) SaveNode(ctx context.Context, traceID string, node *types.Trace
|
|||
|
||||
// LoadNode loads a node from store
|
||||
func (d *Driver) LoadNode(ctx context.Context, traceID string, nodeID string) (*types.TraceNode, error) {
|
||||
// Check if archived and extract if needed
|
||||
archived, err := d.IsArchived(ctx, traceID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to check archive status: %w", err)
|
||||
}
|
||||
if archived {
|
||||
if err := d.unarchive(ctx, traceID); err != nil {
|
||||
return nil, fmt.Errorf("failed to unarchive trace: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
key := d.getKey(traceID, "node", nodeID)
|
||||
|
||||
value, ok := d.store.Get(key)
|
||||
|
|
@ -203,6 +242,15 @@ func (d *Driver) LoadTrace(ctx context.Context, traceID string) (*types.TraceNod
|
|||
|
||||
// SaveSpace persists a space to store
|
||||
func (d *Driver) SaveSpace(ctx context.Context, traceID string, space *types.TraceSpace) error {
|
||||
// Check if archived - archived traces are read-only
|
||||
archived, err := d.IsArchived(ctx, traceID)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to check archive status: %w", err)
|
||||
}
|
||||
if archived {
|
||||
return fmt.Errorf("cannot save space: trace %s is archived (read-only)", traceID)
|
||||
}
|
||||
|
||||
key := d.getKey(traceID, "space", space.ID)
|
||||
|
||||
data, err := json.Marshal(space)
|
||||
|
|
@ -219,6 +267,17 @@ func (d *Driver) SaveSpace(ctx context.Context, traceID string, space *types.Tra
|
|||
|
||||
// LoadSpace loads a space from store
|
||||
func (d *Driver) LoadSpace(ctx context.Context, traceID string, spaceID string) (*types.TraceSpace, error) {
|
||||
// Check if archived and extract if needed
|
||||
archived, err := d.IsArchived(ctx, traceID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to check archive status: %w", err)
|
||||
}
|
||||
if archived {
|
||||
if err := d.unarchive(ctx, traceID); err != nil {
|
||||
return nil, fmt.Errorf("failed to unarchive trace: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
key := d.getKey(traceID, "space", spaceID)
|
||||
|
||||
value, ok := d.store.Get(key)
|
||||
|
|
@ -399,6 +458,15 @@ func (d *Driver) ListSpaceKeys(ctx context.Context, traceID, spaceID string) ([]
|
|||
|
||||
// SaveLog appends a log entry to store
|
||||
func (d *Driver) SaveLog(ctx context.Context, traceID string, log *types.TraceLog) error {
|
||||
// Check if archived - archived traces are read-only
|
||||
archived, err := d.IsArchived(ctx, traceID)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to check archive status: %w", err)
|
||||
}
|
||||
if archived {
|
||||
return fmt.Errorf("cannot save log: trace %s is archived (read-only)", traceID)
|
||||
}
|
||||
|
||||
// Store logs using ArraySlice approach (store as array in a key)
|
||||
key := d.getKey(traceID, "logs", log.NodeID)
|
||||
|
||||
|
|
@ -418,6 +486,17 @@ func (d *Driver) SaveLog(ctx context.Context, traceID string, log *types.TraceLo
|
|||
|
||||
// LoadLogs loads all logs for a trace or specific node from store
|
||||
func (d *Driver) LoadLogs(ctx context.Context, traceID string, nodeID string) ([]*types.TraceLog, error) {
|
||||
// Check if archived and extract if needed
|
||||
archived, err := d.IsArchived(ctx, traceID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to check archive status: %w", err)
|
||||
}
|
||||
if archived {
|
||||
if err := d.unarchive(ctx, traceID); err != nil {
|
||||
return nil, fmt.Errorf("failed to unarchive trace: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
var logs []*types.TraceLog
|
||||
|
||||
if nodeID != "" {
|
||||
|
|
@ -476,6 +555,15 @@ func (d *Driver) loadLogsFromKey(key string) ([]*types.TraceLog, error) {
|
|||
|
||||
// SaveTraceInfo persists trace metadata to store
|
||||
func (d *Driver) SaveTraceInfo(ctx context.Context, info *types.TraceInfo) error {
|
||||
// Allow saving trace info even if archived (for updating archive status)
|
||||
if info.Archived {
|
||||
// If already archived, only allow updating archive-related fields
|
||||
existing, err := d.LoadTraceInfo(ctx, info.ID)
|
||||
if err == nil && existing != nil && existing.Archived && !info.Archived {
|
||||
return fmt.Errorf("cannot unarchive trace: trace %s is archived", info.ID)
|
||||
}
|
||||
}
|
||||
|
||||
key := d.getKey(info.ID, "info")
|
||||
|
||||
data, err := json.Marshal(info)
|
||||
|
|
@ -490,8 +578,8 @@ func (d *Driver) SaveTraceInfo(ctx context.Context, info *types.TraceInfo) error
|
|||
return nil
|
||||
}
|
||||
|
||||
// LoadTraceInfo loads trace metadata from store
|
||||
func (d *Driver) LoadTraceInfo(ctx context.Context, traceID string) (*types.TraceInfo, error) {
|
||||
// loadTraceInfoDirect loads trace info without unarchiving (internal use)
|
||||
func (d *Driver) loadTraceInfoDirect(traceID string) (*types.TraceInfo, error) {
|
||||
key := d.getKey(traceID, "info")
|
||||
|
||||
value, ok := d.store.Get(key)
|
||||
|
|
@ -512,6 +600,22 @@ func (d *Driver) LoadTraceInfo(ctx context.Context, traceID string) (*types.Trac
|
|||
return &info, nil
|
||||
}
|
||||
|
||||
// LoadTraceInfo loads trace metadata from store
|
||||
func (d *Driver) LoadTraceInfo(ctx context.Context, traceID string) (*types.TraceInfo, error) {
|
||||
// Check if archived and extract if needed
|
||||
archived, err := d.IsArchived(ctx, traceID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to check archive status: %w", err)
|
||||
}
|
||||
if archived {
|
||||
if err := d.unarchive(ctx, traceID); err != nil {
|
||||
return nil, fmt.Errorf("failed to unarchive trace: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
return d.loadTraceInfoDirect(traceID)
|
||||
}
|
||||
|
||||
// DeleteTrace removes entire trace from store
|
||||
func (d *Driver) DeleteTrace(ctx context.Context, traceID string) error {
|
||||
// Get all keys
|
||||
|
|
@ -530,6 +634,15 @@ func (d *Driver) DeleteTrace(ctx context.Context, traceID string) error {
|
|||
|
||||
// SaveUpdate persists a trace update event to store (append to list)
|
||||
func (d *Driver) SaveUpdate(ctx context.Context, traceID string, update *types.TraceUpdate) error {
|
||||
// Check if archived - archived traces are read-only
|
||||
archived, err := d.IsArchived(ctx, traceID)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to check archive status: %w", err)
|
||||
}
|
||||
if archived {
|
||||
return fmt.Errorf("cannot save update: trace %s is archived (read-only)", traceID)
|
||||
}
|
||||
|
||||
key := d.getKey(traceID, "updates")
|
||||
|
||||
// Lock to prevent concurrent updates
|
||||
|
|
@ -589,6 +702,185 @@ func (d *Driver) LoadUpdates(ctx context.Context, traceID string, since int64) (
|
|||
}
|
||||
|
||||
// Close closes the store driver
|
||||
// Archive archives a trace by compressing and merging keys
|
||||
func (d *Driver) Archive(ctx context.Context, traceID string) error {
|
||||
// Check if already archived
|
||||
archived, err := d.IsArchived(ctx, traceID)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to check archive status: %w", err)
|
||||
}
|
||||
if archived {
|
||||
return fmt.Errorf("trace %s is already archived", traceID)
|
||||
}
|
||||
|
||||
// Step 1: Update trace info to mark as archived BEFORE creating archive
|
||||
info, err := d.loadTraceInfoDirect(traceID)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to load trace info: %w", err)
|
||||
}
|
||||
if info != nil {
|
||||
now := time.Now().UnixMilli()
|
||||
info.Archived = true
|
||||
info.ArchivedAt = &now
|
||||
// Save trace info directly without archive check
|
||||
key := d.getKey(traceID, "info")
|
||||
infoData, err := json.Marshal(info)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to marshal trace info: %w", err)
|
||||
}
|
||||
if err := d.store.Set(key, string(infoData), 0); err != nil {
|
||||
return fmt.Errorf("failed to save trace info: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
// Step 2: Collect all keys for this trace
|
||||
prefix := d.getKeyPrefix(traceID)
|
||||
allKeys := []string{
|
||||
d.getTraceInfoKey(traceID),
|
||||
d.getUpdatesKey(traceID),
|
||||
}
|
||||
|
||||
// Get all node keys
|
||||
nodePrefix := prefix + "nodes:"
|
||||
nodeKeys, err := d.listKeysByPrefix(ctx, nodePrefix)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to list node keys: %w", err)
|
||||
}
|
||||
allKeys = append(allKeys, nodeKeys...)
|
||||
|
||||
// Get all space keys
|
||||
spacePrefix := prefix + "spaces:"
|
||||
spaceKeys, err := d.listKeysByPrefix(ctx, spacePrefix)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to list space keys: %w", err)
|
||||
}
|
||||
allKeys = append(allKeys, spaceKeys...)
|
||||
|
||||
// Get all log keys
|
||||
logPrefix := prefix + "logs:"
|
||||
logKeys, err := d.listKeysByPrefix(ctx, logPrefix)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to list log keys: %w", err)
|
||||
}
|
||||
allKeys = append(allKeys, logKeys...)
|
||||
|
||||
// Step 3: Collect all data into a single map
|
||||
archiveData := make(map[string]json.RawMessage)
|
||||
for _, key := range allKeys {
|
||||
data, ok := d.store.Get(key)
|
||||
if !ok {
|
||||
continue // Skip missing keys
|
||||
}
|
||||
// Convert to string then to bytes
|
||||
dataStr, ok := data.(string)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
archiveData[key] = json.RawMessage(dataStr)
|
||||
}
|
||||
|
||||
// Step 4: Marshal to JSON
|
||||
jsonData, err := json.Marshal(archiveData)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to marshal archive data: %w", err)
|
||||
}
|
||||
|
||||
// Step 5: Compress with gzip
|
||||
var compressedBuf bytes.Buffer
|
||||
gzipWriter := gzip.NewWriter(&compressedBuf)
|
||||
if _, err := gzipWriter.Write(jsonData); err != nil {
|
||||
return fmt.Errorf("failed to compress archive: %w", err)
|
||||
}
|
||||
if err := gzipWriter.Close(); err != nil {
|
||||
return fmt.Errorf("failed to close gzip writer: %w", err)
|
||||
}
|
||||
|
||||
// Step 6: Save compressed archive
|
||||
archiveKey := d.getArchiveKey(traceID)
|
||||
if err := d.store.Set(archiveKey, compressedBuf.String(), 0); err != nil {
|
||||
return fmt.Errorf("failed to save archive: %w", err)
|
||||
}
|
||||
|
||||
// Step 7: Delete original keys (except trace info and archive)
|
||||
for _, key := range allKeys {
|
||||
if key == d.getTraceInfoKey(traceID) {
|
||||
continue // Keep trace info
|
||||
}
|
||||
_ = d.store.Del(key) // Ignore errors on delete
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// IsArchived checks if a trace is archived
|
||||
func (d *Driver) IsArchived(ctx context.Context, traceID string) (bool, error) {
|
||||
archiveKey := d.getArchiveKey(traceID)
|
||||
exists := d.store.Has(archiveKey)
|
||||
return exists, nil
|
||||
}
|
||||
|
||||
// unarchive extracts an archived trace (helper method)
|
||||
func (d *Driver) unarchive(ctx context.Context, traceID string) error {
|
||||
archiveKey := d.getArchiveKey(traceID)
|
||||
|
||||
// Get compressed archive
|
||||
compressedData, ok := d.store.Get(archiveKey)
|
||||
if !ok {
|
||||
return fmt.Errorf("archive not found for trace: %s", traceID)
|
||||
}
|
||||
|
||||
// Convert to string then to bytes
|
||||
compressedStr, ok := compressedData.(string)
|
||||
if !ok {
|
||||
return fmt.Errorf("invalid archive data type")
|
||||
}
|
||||
|
||||
// Decompress
|
||||
gzipReader, err := gzip.NewReader(bytes.NewReader([]byte(compressedStr)))
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to create gzip reader: %w", err)
|
||||
}
|
||||
defer gzipReader.Close()
|
||||
|
||||
jsonData, err := io.ReadAll(gzipReader)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to decompress archive: %w", err)
|
||||
}
|
||||
|
||||
// Unmarshal archive data
|
||||
var archiveData map[string]json.RawMessage
|
||||
if err := json.Unmarshal(jsonData, &archiveData); err != nil {
|
||||
return fmt.Errorf("failed to unmarshal archive: %w", err)
|
||||
}
|
||||
|
||||
// Restore all keys
|
||||
for key, value := range archiveData {
|
||||
if err := d.store.Set(key, string(value), 0); err != nil {
|
||||
return fmt.Errorf("failed to restore key %s: %w", key, err)
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// listKeysByPrefix lists all keys with a given prefix (helper method)
|
||||
func (d *Driver) listKeysByPrefix(ctx context.Context, prefix string) ([]string, error) {
|
||||
// Get all keys from store and filter by prefix
|
||||
allKeys := d.store.Keys()
|
||||
var matchingKeys []string
|
||||
for _, key := range allKeys {
|
||||
if strings.HasPrefix(key, prefix) {
|
||||
matchingKeys = append(matchingKeys, key)
|
||||
}
|
||||
}
|
||||
return matchingKeys, nil
|
||||
}
|
||||
|
||||
// getArchiveKey returns the store key for an archived trace
|
||||
func (d *Driver) getArchiveKey(traceID string) string {
|
||||
return fmt.Sprintf("trace:%s:archive", traceID)
|
||||
}
|
||||
|
||||
func (d *Driver) Close() error {
|
||||
// Store connection is managed by gou, no cleanup needed
|
||||
return nil
|
||||
|
|
|
|||
|
|
@ -138,7 +138,7 @@ func New(ctx context.Context, driver string, option *types.TraceOption, driverOp
|
|||
}
|
||||
|
||||
// Create Manager instance with the driver
|
||||
manager, err := NewManager(ctx, traceID, drv)
|
||||
manager, err := NewManager(ctx, traceID, drv, option)
|
||||
if err != nil {
|
||||
return "", nil, fmt.Errorf("failed to create manager: %w", err)
|
||||
}
|
||||
|
|
@ -227,7 +227,8 @@ func LoadFromStorage(ctx context.Context, driver string, traceID string, driverO
|
|||
// Create Manager instance with the driver
|
||||
// Note: We need to reconstruct the manager from stored data
|
||||
// TODO: Implement proper restoration of manager state from storage
|
||||
manager, err := NewManager(ctx, traceID, drv)
|
||||
// For loaded traces, we don't have the original option, so pass nil
|
||||
manager, err := NewManager(ctx, traceID, drv, nil)
|
||||
if err != nil {
|
||||
drv.Close()
|
||||
return "", nil, fmt.Errorf("failed to create manager: %w", err)
|
||||
|
|
|
|||
92
trace/trace_archive_test.go
Normal file
92
trace/trace_archive_test.go
Normal file
|
|
@ -0,0 +1,92 @@
|
|||
package trace_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/yaoapp/yao/trace"
|
||||
"github.com/yaoapp/yao/trace/types"
|
||||
)
|
||||
|
||||
func TestArchive(t *testing.T) {
|
||||
drivers := trace.GetTestDrivers()
|
||||
|
||||
for _, d := range drivers {
|
||||
t.Run(d.Name, func(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
|
||||
// Create a trace with AutoArchive disabled
|
||||
traceID, manager, err := trace.New(ctx, d.DriverType, &types.TraceOption{
|
||||
AutoArchive: false,
|
||||
}, d.DriverOptions...)
|
||||
assert.NoError(t, err)
|
||||
defer trace.Release(traceID)
|
||||
defer trace.Remove(ctx, d.DriverType, traceID, d.DriverOptions...)
|
||||
|
||||
// Create some test data
|
||||
_, err = manager.Add("test input", types.TraceNodeOption{
|
||||
Label: "Test Node",
|
||||
})
|
||||
assert.NoError(t, err)
|
||||
|
||||
manager.Info("Test log message", map[string]any{
|
||||
"key": "value",
|
||||
})
|
||||
|
||||
err = manager.Complete()
|
||||
assert.NoError(t, err)
|
||||
|
||||
err = manager.MarkComplete()
|
||||
assert.NoError(t, err)
|
||||
|
||||
// Wait a bit for completion
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
|
||||
// Note: Archive functionality is tested at the driver level
|
||||
// Here we just verify that traces can be created and completed
|
||||
// without AutoArchive enabled
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAutoArchive(t *testing.T) {
|
||||
drivers := trace.GetTestDrivers()
|
||||
|
||||
for _, d := range drivers {
|
||||
t.Run(d.Name, func(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
|
||||
// Create a trace with AutoArchive enabled
|
||||
traceID, manager, err := trace.New(ctx, d.DriverType, &types.TraceOption{
|
||||
AutoArchive: true,
|
||||
}, d.DriverOptions...)
|
||||
assert.NoError(t, err)
|
||||
defer trace.Release(traceID)
|
||||
defer trace.Remove(ctx, d.DriverType, traceID, d.DriverOptions...)
|
||||
|
||||
// Create some test data
|
||||
_, err = manager.Add("test input", types.TraceNodeOption{
|
||||
Label: "Test Node",
|
||||
})
|
||||
assert.NoError(t, err)
|
||||
|
||||
manager.Info("Test log message", map[string]any{
|
||||
"key": "value",
|
||||
})
|
||||
|
||||
err = manager.Complete()
|
||||
assert.NoError(t, err)
|
||||
|
||||
err = manager.MarkComplete()
|
||||
assert.NoError(t, err)
|
||||
|
||||
// Wait for auto-archive to complete
|
||||
time.Sleep(200 * time.Millisecond)
|
||||
|
||||
// Trace should still be accessible after auto-archive
|
||||
assert.True(t, trace.IsLoaded(traceID))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
|
@ -75,6 +75,12 @@ type Driver interface {
|
|||
// LoadUpdates loads trace update events from storage (filtering by timestamp in milliseconds)
|
||||
LoadUpdates(ctx context.Context, traceID string, since int64) ([]*TraceUpdate, error)
|
||||
|
||||
// Archive archives a trace (compress and make read-only)
|
||||
Archive(ctx context.Context, traceID string) error
|
||||
|
||||
// IsArchived checks if a trace is archived
|
||||
IsArchived(ctx context.Context, traceID string) (bool, error)
|
||||
|
||||
// Close closes the driver and releases resources
|
||||
Close() error
|
||||
}
|
||||
|
|
|
|||
|
|
@ -200,6 +200,8 @@ type TraceInfo struct {
|
|||
Manager Manager `json:"-"` // Not persisted
|
||||
CreatedAt int64 `json:"created_at"` // milliseconds since epoch
|
||||
UpdatedAt int64 `json:"updated_at"` // milliseconds since epoch
|
||||
ArchivedAt *int64 `json:"archived_at,omitempty"` // milliseconds since epoch, nil if not archived
|
||||
Archived bool `json:"archived"` // Whether this trace is archived (read-only)
|
||||
CreatedBy string `json:"__yao_created_by,omitempty"`
|
||||
UpdatedBy string `json:"__yao_updated_by,omitempty"`
|
||||
TeamID string `json:"__yao_team_id,omitempty"`
|
||||
|
|
@ -214,4 +216,7 @@ type TraceOption struct {
|
|||
TeamID string // Team ID
|
||||
TenantID string // Tenant ID
|
||||
Metadata map[string]any // Additional metadata
|
||||
AutoArchive bool // Automatically archive when trace completes/fails
|
||||
ArchiveOnClose bool // Archive on explicit Close() call
|
||||
ArchiveCompressLevel int // gzip compression level (0-9, default: gzip.DefaultCompression)
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue