Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
9ed6c2638d | ||
|
|
b2cebe1ba7 | ||
|
|
d8fe3f5a3c | ||
|
|
87ceea3d30 | ||
|
|
8baf16a9d3 | ||
|
|
b77c7b5f56 | ||
|
|
9315fa6d32 |
@@ -0,0 +1,331 @@
|
||||
package batching
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// GitOp represents a git operation to be batched
|
||||
type GitOp struct {
|
||||
OpType string // "commit", "push", "merge"
|
||||
Branch string
|
||||
Message string
|
||||
Files []string
|
||||
Timestamp time.Time
|
||||
ID string
|
||||
}
|
||||
|
||||
// GitBatch represents a batch of git operations
|
||||
type GitBatch struct {
|
||||
ID string
|
||||
Operations []*GitOp
|
||||
CreatedAt time.Time
|
||||
ExecutedAt time.Time
|
||||
Status string // "pending", "executing", "completed", "failed"
|
||||
Error error
|
||||
}
|
||||
|
||||
// GitBatcher batches git operations for efficient execution
|
||||
type GitBatcher struct {
|
||||
mu sync.RWMutex
|
||||
queue []*GitOp
|
||||
maxBatchSize int
|
||||
maxBatchAge time.Duration
|
||||
lastFlushTime time.Time
|
||||
executedBatches []*GitBatch
|
||||
pendingBatches []*GitBatch
|
||||
stats *BatchStats
|
||||
flushChan chan struct{}
|
||||
stopChan chan struct{}
|
||||
}
|
||||
|
||||
// BatchStats tracks batching statistics
|
||||
type BatchStats struct {
|
||||
TotalOps int
|
||||
TotalBatches int
|
||||
AvgOpsPerBatch float64
|
||||
NetworkSavings int // Estimated network round trips saved
|
||||
TotalExecuteTime time.Duration
|
||||
}
|
||||
|
||||
// NewGitBatcher creates a new git batcher
|
||||
func NewGitBatcher(maxBatchSize int, maxBatchAge time.Duration) *GitBatcher {
|
||||
if maxBatchSize <= 0 {
|
||||
maxBatchSize = 10
|
||||
}
|
||||
if maxBatchAge <= 0 {
|
||||
maxBatchAge = 5 * time.Second
|
||||
}
|
||||
|
||||
return &GitBatcher{
|
||||
queue: make([]*GitOp, 0),
|
||||
maxBatchSize: maxBatchSize,
|
||||
maxBatchAge: maxBatchAge,
|
||||
lastFlushTime: time.Now(),
|
||||
executedBatches: make([]*GitBatch, 0),
|
||||
pendingBatches: make([]*GitBatch, 0),
|
||||
stats: &BatchStats{
|
||||
TotalOps: 0,
|
||||
TotalBatches: 0,
|
||||
},
|
||||
flushChan: make(chan struct{}, 1),
|
||||
stopChan: make(chan struct{}),
|
||||
}
|
||||
}
|
||||
|
||||
// Enqueue adds a git operation to the queue
|
||||
func (gb *GitBatcher) Enqueue(op *GitOp) {
|
||||
if op == nil {
|
||||
return
|
||||
}
|
||||
|
||||
op.Timestamp = time.Now()
|
||||
|
||||
gb.mu.Lock()
|
||||
defer gb.mu.Unlock()
|
||||
|
||||
gb.queue = append(gb.queue, op)
|
||||
gb.stats.TotalOps++
|
||||
|
||||
// Auto-flush if batch is full
|
||||
if len(gb.queue) >= gb.maxBatchSize {
|
||||
gb.flushLocked()
|
||||
}
|
||||
}
|
||||
|
||||
// flushLocked creates a batch from queued operations (must be called with lock held)
|
||||
func (gb *GitBatcher) flushLocked() {
|
||||
if len(gb.queue) == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
batch := &GitBatch{
|
||||
ID: fmt.Sprintf("batch-%d", gb.stats.TotalBatches),
|
||||
Operations: make([]*GitOp, len(gb.queue)),
|
||||
CreatedAt: time.Now(),
|
||||
Status: "pending",
|
||||
}
|
||||
|
||||
copy(batch.Operations, gb.queue)
|
||||
|
||||
gb.pendingBatches = append(gb.pendingBatches, batch)
|
||||
gb.queue = make([]*GitOp, 0)
|
||||
gb.lastFlushTime = time.Now()
|
||||
gb.stats.TotalBatches++
|
||||
}
|
||||
|
||||
// Flush manually flushes the current batch
|
||||
func (gb *GitBatcher) Flush() {
|
||||
gb.mu.Lock()
|
||||
defer gb.mu.Unlock()
|
||||
|
||||
gb.flushLocked()
|
||||
}
|
||||
|
||||
// GetPendingBatch returns the next pending batch without removing it
|
||||
func (gb *GitBatcher) GetPendingBatch() *GitBatch {
|
||||
gb.mu.RLock()
|
||||
defer gb.mu.RUnlock()
|
||||
|
||||
if len(gb.pendingBatches) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
return gb.pendingBatches[0]
|
||||
}
|
||||
|
||||
// MarkBatchExecuting marks a batch as executing
|
||||
func (gb *GitBatcher) MarkBatchExecuting(batchID string) {
|
||||
gb.mu.Lock()
|
||||
defer gb.mu.Unlock()
|
||||
|
||||
for _, batch := range gb.pendingBatches {
|
||||
if batch.ID == batchID {
|
||||
batch.Status = "executing"
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// MarkBatchCompleted marks a batch as completed and removes from pending
|
||||
func (gb *GitBatcher) MarkBatchCompleted(batchID string) {
|
||||
gb.mu.Lock()
|
||||
defer gb.mu.Unlock()
|
||||
|
||||
var idx int
|
||||
var found *GitBatch
|
||||
for i, batch := range gb.pendingBatches {
|
||||
if batch.ID == batchID {
|
||||
idx = i
|
||||
found = batch
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if found != nil {
|
||||
found.Status = "completed"
|
||||
found.ExecutedAt = time.Now()
|
||||
|
||||
// Move to executed batches
|
||||
gb.executedBatches = append(gb.executedBatches, found)
|
||||
|
||||
// Remove from pending
|
||||
gb.pendingBatches = append(gb.pendingBatches[:idx], gb.pendingBatches[idx+1:]...)
|
||||
}
|
||||
}
|
||||
|
||||
// MarkBatchFailed marks a batch as failed with an error
|
||||
func (gb *GitBatcher) MarkBatchFailed(batchID string, err error) {
|
||||
gb.mu.Lock()
|
||||
defer gb.mu.Unlock()
|
||||
|
||||
var found *GitBatch
|
||||
for _, batch := range gb.pendingBatches {
|
||||
if batch.ID == batchID {
|
||||
found = batch
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if found != nil {
|
||||
found.Status = "failed"
|
||||
found.Error = err
|
||||
found.ExecutedAt = time.Now()
|
||||
|
||||
// Keep in pending (for retry logic)
|
||||
// Could also move to failed queue
|
||||
}
|
||||
}
|
||||
|
||||
// QueueSize returns the current queue size
|
||||
func (gb *GitBatcher) QueueSize() int {
|
||||
gb.mu.RLock()
|
||||
defer gb.mu.RUnlock()
|
||||
|
||||
return len(gb.queue)
|
||||
}
|
||||
|
||||
// PendingBatchCount returns the number of pending batches
|
||||
func (gb *GitBatcher) PendingBatchCount() int {
|
||||
gb.mu.RLock()
|
||||
defer gb.mu.RUnlock()
|
||||
|
||||
return len(gb.pendingBatches)
|
||||
}
|
||||
|
||||
// GetStats returns batching statistics
|
||||
func (gb *GitBatcher) GetStats() *BatchStats {
|
||||
gb.mu.RLock()
|
||||
defer gb.mu.RUnlock()
|
||||
|
||||
stats := *gb.stats
|
||||
if stats.TotalBatches > 0 {
|
||||
stats.AvgOpsPerBatch = float64(stats.TotalOps) / float64(stats.TotalBatches)
|
||||
// Estimated savings: each batch saves (ops-1) round trips
|
||||
stats.NetworkSavings = stats.TotalOps - stats.TotalBatches
|
||||
}
|
||||
|
||||
return &stats
|
||||
}
|
||||
|
||||
// GetExecutedBatches returns all executed batches
|
||||
func (gb *GitBatcher) GetExecutedBatches() []*GitBatch {
|
||||
gb.mu.RLock()
|
||||
defer gb.mu.RUnlock()
|
||||
|
||||
result := make([]*GitBatch, len(gb.executedBatches))
|
||||
copy(result, gb.executedBatches)
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
// GetBatchByID returns a specific batch by ID
|
||||
func (gb *GitBatcher) GetBatchByID(batchID string) *GitBatch {
|
||||
gb.mu.RLock()
|
||||
defer gb.mu.RUnlock()
|
||||
|
||||
for _, batch := range gb.pendingBatches {
|
||||
if batch.ID == batchID {
|
||||
return batch
|
||||
}
|
||||
}
|
||||
|
||||
for _, batch := range gb.executedBatches {
|
||||
if batch.ID == batchID {
|
||||
return batch
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// TimeSinceLastFlush returns time since last flush
|
||||
func (gb *GitBatcher) TimeSinceLastFlush() time.Duration {
|
||||
gb.mu.RLock()
|
||||
defer gb.mu.RUnlock()
|
||||
|
||||
return time.Since(gb.lastFlushTime)
|
||||
}
|
||||
|
||||
// ShouldFlush checks if batch should be flushed based on age
|
||||
func (gb *GitBatcher) ShouldFlush() bool {
|
||||
gb.mu.RLock()
|
||||
defer gb.mu.RUnlock()
|
||||
|
||||
if len(gb.queue) == 0 {
|
||||
return false
|
||||
}
|
||||
|
||||
return time.Since(gb.lastFlushTime) >= gb.maxBatchAge
|
||||
}
|
||||
|
||||
// Clear clears all pending operations and batches
|
||||
func (gb *GitBatcher) Clear() {
|
||||
gb.mu.Lock()
|
||||
defer gb.mu.Unlock()
|
||||
|
||||
gb.queue = make([]*GitOp, 0)
|
||||
gb.pendingBatches = make([]*GitBatch, 0)
|
||||
gb.executedBatches = make([]*GitBatch, 0)
|
||||
}
|
||||
|
||||
// GetQueuedOps returns a copy of queued operations
|
||||
func (gb *GitBatcher) GetQueuedOps() []*GitOp {
|
||||
gb.mu.RLock()
|
||||
defer gb.mu.RUnlock()
|
||||
|
||||
ops := make([]*GitOp, len(gb.queue))
|
||||
copy(ops, gb.queue)
|
||||
|
||||
return ops
|
||||
}
|
||||
|
||||
// CalculateNetworkSavings calculates estimated network round trips saved
|
||||
func (gb *GitBatcher) CalculateNetworkSavings() int {
|
||||
gb.mu.RLock()
|
||||
defer gb.mu.RUnlock()
|
||||
|
||||
totalSavings := 0
|
||||
// Each batch of N operations saves N-1 round trips
|
||||
for _, batch := range gb.executedBatches {
|
||||
if len(batch.Operations) > 1 {
|
||||
totalSavings += len(batch.Operations) - 1
|
||||
}
|
||||
}
|
||||
|
||||
return totalSavings
|
||||
}
|
||||
|
||||
// GetBatchInfo returns human-readable batch information
|
||||
func (batch *GitBatch) GetInfo() map[string]interface{} {
|
||||
return map[string]interface{}{
|
||||
"id": batch.ID,
|
||||
"status": batch.Status,
|
||||
"op_count": len(batch.Operations),
|
||||
"created_at": batch.CreatedAt,
|
||||
"executed_at": batch.ExecutedAt,
|
||||
"duration": batch.ExecutedAt.Sub(batch.CreatedAt),
|
||||
"error": batch.Error,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,355 @@
|
||||
package batching
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
func TestNewGitBatcher(t *testing.T) {
|
||||
batcher := NewGitBatcher(10, 5*time.Second)
|
||||
assert.NotNil(t, batcher)
|
||||
assert.Equal(t, 0, batcher.QueueSize())
|
||||
}
|
||||
|
||||
func TestEnqueueOperation(t *testing.T) {
|
||||
batcher := NewGitBatcher(10, 5*time.Second)
|
||||
|
||||
op := &GitOp{
|
||||
OpType: "commit",
|
||||
Branch: "main",
|
||||
Message: "Add feature",
|
||||
Files: []string{"file1.go"},
|
||||
}
|
||||
|
||||
batcher.Enqueue(op)
|
||||
assert.Equal(t, 1, batcher.QueueSize())
|
||||
}
|
||||
|
||||
func TestEnqueueMultipleOps(t *testing.T) {
|
||||
batcher := NewGitBatcher(10, 5*time.Second)
|
||||
|
||||
for i := 0; i < 5; i++ {
|
||||
op := &GitOp{
|
||||
OpType: "commit",
|
||||
Branch: "main",
|
||||
Message: "Commit",
|
||||
}
|
||||
batcher.Enqueue(op)
|
||||
}
|
||||
|
||||
assert.Equal(t, 5, batcher.QueueSize())
|
||||
}
|
||||
|
||||
func TestAutoFlushOnMaxBatchSize(t *testing.T) {
|
||||
batcher := NewGitBatcher(5, 10*time.Second)
|
||||
|
||||
for i := 0; i < 5; i++ {
|
||||
op := &GitOp{
|
||||
OpType: "commit",
|
||||
Branch: "main",
|
||||
Message: "Commit",
|
||||
}
|
||||
batcher.Enqueue(op)
|
||||
}
|
||||
|
||||
// After 5 ops, should auto-flush
|
||||
assert.Equal(t, 0, batcher.QueueSize())
|
||||
assert.Equal(t, 1, batcher.PendingBatchCount())
|
||||
}
|
||||
|
||||
func TestManualFlush(t *testing.T) {
|
||||
batcher := NewGitBatcher(10, 5*time.Second)
|
||||
|
||||
op := &GitOp{
|
||||
OpType: "commit",
|
||||
Branch: "main",
|
||||
Message: "Commit",
|
||||
}
|
||||
batcher.Enqueue(op)
|
||||
assert.Equal(t, 1, batcher.QueueSize())
|
||||
|
||||
batcher.Flush()
|
||||
assert.Equal(t, 0, batcher.QueueSize())
|
||||
assert.Equal(t, 1, batcher.PendingBatchCount())
|
||||
}
|
||||
|
||||
func TestGetPendingBatch(t *testing.T) {
|
||||
batcher := NewGitBatcher(10, 5*time.Second)
|
||||
|
||||
op := &GitOp{
|
||||
OpType: "commit",
|
||||
Branch: "main",
|
||||
Message: "Commit",
|
||||
}
|
||||
batcher.Enqueue(op)
|
||||
batcher.Flush()
|
||||
|
||||
batch := batcher.GetPendingBatch()
|
||||
assert.NotNil(t, batch)
|
||||
assert.Equal(t, 1, len(batch.Operations))
|
||||
}
|
||||
|
||||
func TestMarkBatchExecuting(t *testing.T) {
|
||||
batcher := NewGitBatcher(10, 5*time.Second)
|
||||
|
||||
op := &GitOp{OpType: "commit"}
|
||||
batcher.Enqueue(op)
|
||||
batcher.Flush()
|
||||
|
||||
batch := batcher.GetPendingBatch()
|
||||
batcher.MarkBatchExecuting(batch.ID)
|
||||
|
||||
updated := batcher.GetBatchByID(batch.ID)
|
||||
assert.Equal(t, "executing", updated.Status)
|
||||
}
|
||||
|
||||
func TestMarkBatchCompleted(t *testing.T) {
|
||||
batcher := NewGitBatcher(10, 5*time.Second)
|
||||
|
||||
op := &GitOp{OpType: "commit"}
|
||||
batcher.Enqueue(op)
|
||||
batcher.Flush()
|
||||
|
||||
batch := batcher.GetPendingBatch()
|
||||
batcher.MarkBatchCompleted(batch.ID)
|
||||
|
||||
executed := batcher.GetExecutedBatches()
|
||||
assert.Equal(t, 1, len(executed))
|
||||
assert.Equal(t, "completed", executed[0].Status)
|
||||
}
|
||||
|
||||
func TestMarkBatchFailed(t *testing.T) {
|
||||
batcher := NewGitBatcher(10, 5*time.Second)
|
||||
|
||||
op := &GitOp{OpType: "commit"}
|
||||
batcher.Enqueue(op)
|
||||
batcher.Flush()
|
||||
|
||||
batch := batcher.GetPendingBatch()
|
||||
testErr := assert.AnError
|
||||
batcher.MarkBatchFailed(batch.ID, testErr)
|
||||
|
||||
failed := batcher.GetBatchByID(batch.ID)
|
||||
assert.Equal(t, "failed", failed.Status)
|
||||
assert.Error(t, failed.Error)
|
||||
}
|
||||
|
||||
func TestGetStats(t *testing.T) {
|
||||
batcher := NewGitBatcher(5, 5*time.Second)
|
||||
|
||||
// Add 10 ops (will create 2 batches of 5 each)
|
||||
for i := 0; i < 10; i++ {
|
||||
op := &GitOp{OpType: "commit"}
|
||||
batcher.Enqueue(op)
|
||||
}
|
||||
|
||||
stats := batcher.GetStats()
|
||||
assert.Equal(t, 10, stats.TotalOps)
|
||||
assert.Equal(t, 2, stats.TotalBatches)
|
||||
assert.Equal(t, 5.0, stats.AvgOpsPerBatch)
|
||||
// 10 ops in 2 batches saves 8 round trips (5-1 + 5-1)
|
||||
assert.Equal(t, 8, stats.NetworkSavings)
|
||||
}
|
||||
|
||||
func TestQueueSize(t *testing.T) {
|
||||
batcher := NewGitBatcher(10, 5*time.Second)
|
||||
|
||||
op := &GitOp{OpType: "commit"}
|
||||
batcher.Enqueue(op)
|
||||
|
||||
assert.Equal(t, 1, batcher.QueueSize())
|
||||
}
|
||||
|
||||
func TestPendingBatchCount(t *testing.T) {
|
||||
batcher := NewGitBatcher(10, 5*time.Second)
|
||||
|
||||
op := &GitOp{OpType: "commit"}
|
||||
batcher.Enqueue(op)
|
||||
batcher.Flush()
|
||||
|
||||
assert.Equal(t, 1, batcher.PendingBatchCount())
|
||||
}
|
||||
|
||||
func TestGetExecutedBatches(t *testing.T) {
|
||||
batcher := NewGitBatcher(10, 5*time.Second)
|
||||
|
||||
// Create and execute batches
|
||||
for i := 0; i < 2; i++ {
|
||||
op := &GitOp{OpType: "commit"}
|
||||
batcher.Enqueue(op)
|
||||
batcher.Flush()
|
||||
|
||||
batch := batcher.GetPendingBatch()
|
||||
batcher.MarkBatchCompleted(batch.ID)
|
||||
}
|
||||
|
||||
executed := batcher.GetExecutedBatches()
|
||||
assert.Equal(t, 2, len(executed))
|
||||
}
|
||||
|
||||
func TestTimeSinceLastFlush(t *testing.T) {
|
||||
batcher := NewGitBatcher(10, 5*time.Second)
|
||||
|
||||
op := &GitOp{OpType: "commit"}
|
||||
batcher.Enqueue(op)
|
||||
batcher.Flush()
|
||||
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
elapsed := batcher.TimeSinceLastFlush()
|
||||
|
||||
assert.Greater(t, elapsed, 50*time.Millisecond)
|
||||
assert.Less(t, elapsed, 200*time.Millisecond)
|
||||
}
|
||||
|
||||
func TestShouldFlush(t *testing.T) {
|
||||
batcher := NewGitBatcher(100, 100*time.Millisecond)
|
||||
|
||||
// Empty queue should not flush
|
||||
assert.False(t, batcher.ShouldFlush())
|
||||
|
||||
// Enqueue but not old enough
|
||||
op := &GitOp{OpType: "commit"}
|
||||
batcher.Enqueue(op)
|
||||
assert.False(t, batcher.ShouldFlush())
|
||||
|
||||
// Wait for age to exceed max age
|
||||
time.Sleep(150 * time.Millisecond)
|
||||
assert.True(t, batcher.ShouldFlush())
|
||||
}
|
||||
|
||||
func TestClear(t *testing.T) {
|
||||
batcher := NewGitBatcher(10, 5*time.Second)
|
||||
|
||||
op := &GitOp{OpType: "commit"}
|
||||
batcher.Enqueue(op)
|
||||
batcher.Flush()
|
||||
|
||||
assert.Equal(t, 1, batcher.PendingBatchCount())
|
||||
|
||||
batcher.Clear()
|
||||
assert.Equal(t, 0, batcher.QueueSize())
|
||||
assert.Equal(t, 0, batcher.PendingBatchCount())
|
||||
}
|
||||
|
||||
func TestGetQueuedOps(t *testing.T) {
|
||||
batcher := NewGitBatcher(10, 5*time.Second)
|
||||
|
||||
ops := []*GitOp{
|
||||
{OpType: "commit", Message: "Commit 1"},
|
||||
{OpType: "commit", Message: "Commit 2"},
|
||||
{OpType: "commit", Message: "Commit 3"},
|
||||
}
|
||||
|
||||
for _, op := range ops {
|
||||
batcher.Enqueue(op)
|
||||
}
|
||||
|
||||
queued := batcher.GetQueuedOps()
|
||||
assert.Equal(t, 3, len(queued))
|
||||
assert.Equal(t, "Commit 1", queued[0].Message)
|
||||
assert.Equal(t, "Commit 3", queued[2].Message)
|
||||
}
|
||||
|
||||
func TestCalculateNetworkSavings(t *testing.T) {
|
||||
batcher := NewGitBatcher(3, 5*time.Second)
|
||||
|
||||
// Add 6 ops (will create 2 batches of 3 each)
|
||||
for i := 0; i < 6; i++ {
|
||||
op := &GitOp{OpType: "commit"}
|
||||
batcher.Enqueue(op)
|
||||
}
|
||||
|
||||
// Mark both batches as completed
|
||||
for i := 0; i < 2; i++ {
|
||||
batch := batcher.GetPendingBatch()
|
||||
if batch != nil {
|
||||
batcher.MarkBatchCompleted(batch.ID)
|
||||
}
|
||||
}
|
||||
|
||||
savings := batcher.CalculateNetworkSavings()
|
||||
// 2 batches of 3 each saves 4 round trips (3-1 + 3-1)
|
||||
assert.Equal(t, 4, savings)
|
||||
}
|
||||
|
||||
func TestGetBatchByID(t *testing.T) {
|
||||
batcher := NewGitBatcher(10, 5*time.Second)
|
||||
|
||||
op := &GitOp{OpType: "commit"}
|
||||
batcher.Enqueue(op)
|
||||
batcher.Flush()
|
||||
|
||||
batch := batcher.GetPendingBatch()
|
||||
retrieved := batcher.GetBatchByID(batch.ID)
|
||||
|
||||
assert.NotNil(t, retrieved)
|
||||
assert.Equal(t, batch.ID, retrieved.ID)
|
||||
}
|
||||
|
||||
func TestGetBatchInfo(t *testing.T) {
|
||||
batch := &GitBatch{
|
||||
ID: "test-batch",
|
||||
Status: "completed",
|
||||
CreatedAt: time.Now(),
|
||||
ExecutedAt: time.Now().Add(1 * time.Second),
|
||||
}
|
||||
|
||||
info := batch.GetInfo()
|
||||
assert.Equal(t, "test-batch", info["id"])
|
||||
assert.Equal(t, "completed", info["status"])
|
||||
}
|
||||
|
||||
func TestMultipleBatches(t *testing.T) {
|
||||
batcher := NewGitBatcher(3, 5*time.Second)
|
||||
|
||||
// Create 3 batches
|
||||
for batch := 0; batch < 3; batch++ {
|
||||
for i := 0; i < 3; i++ {
|
||||
op := &GitOp{
|
||||
OpType: "commit",
|
||||
Branch: "main",
|
||||
}
|
||||
batcher.Enqueue(op)
|
||||
}
|
||||
}
|
||||
|
||||
// All 3 batches should be pending
|
||||
assert.Equal(t, 3, batcher.PendingBatchCount())
|
||||
assert.Equal(t, 0, batcher.QueueSize())
|
||||
}
|
||||
|
||||
func TestEnqueueNil(t *testing.T) {
|
||||
batcher := NewGitBatcher(10, 5*time.Second)
|
||||
|
||||
// Enqueueing nil should not fail
|
||||
batcher.Enqueue(nil)
|
||||
assert.Equal(t, 0, batcher.QueueSize())
|
||||
}
|
||||
|
||||
func BenchmarkEnqueue(b *testing.B) {
|
||||
batcher := NewGitBatcher(1000, 10*time.Second)
|
||||
|
||||
for i := 0; i < b.N; i++ {
|
||||
op := &GitOp{
|
||||
OpType: "commit",
|
||||
Branch: "main",
|
||||
Message: "Commit",
|
||||
}
|
||||
batcher.Enqueue(op)
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkFlush(b *testing.B) {
|
||||
batcher := NewGitBatcher(1000, 10*time.Second)
|
||||
|
||||
for i := 0; i < b.N; i++ {
|
||||
op := &GitOp{OpType: "commit"}
|
||||
batcher.Enqueue(op)
|
||||
|
||||
if (i + 1) % 100 == 0 {
|
||||
batcher.Flush()
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,373 @@
|
||||
package batching
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// LLMRequest represents a single LLM request to be batched
|
||||
type LLMRequest struct {
|
||||
ID string `json:"id"`
|
||||
Type string `json:"type"` // "implementer", "judge", "planner"
|
||||
Model string `json:"model"`
|
||||
Prompt string `json:"prompt"`
|
||||
Metadata map[string]interface{} `json:"metadata,omitempty"`
|
||||
Timestamp time.Time `json:"timestamp"`
|
||||
ResultCh chan *LLMResult `json:"-"`
|
||||
}
|
||||
|
||||
// LLMResult represents the result of a single LLM request
|
||||
type LLMResult struct {
|
||||
RequestID string `json:"request_id"`
|
||||
Response string `json:"response"`
|
||||
Error error `json:"error,omitempty"`
|
||||
Duration time.Duration `json:"duration"`
|
||||
TokenCount int `json:"token_count"`
|
||||
Metadata map[string]interface{} `json:"metadata,omitempty"`
|
||||
}
|
||||
|
||||
// LLMBatch represents a batch of LLM requests
|
||||
type LLMBatch struct {
|
||||
ID string
|
||||
Requests []*LLMRequest
|
||||
Model string
|
||||
Type string
|
||||
CreatedAt time.Time
|
||||
ExecutedAt time.Time
|
||||
Status string // "pending", "executing", "completed", "failed"
|
||||
Error error
|
||||
Results map[string]*LLMResult
|
||||
ExecutionTime time.Duration
|
||||
}
|
||||
|
||||
// LLMBatcher batches LLM requests for efficient API usage
|
||||
type LLMBatcher struct {
|
||||
mu sync.RWMutex
|
||||
queue []*LLMRequest
|
||||
maxBatchSize int
|
||||
maxBatchAge time.Duration
|
||||
lastFlushTime time.Time
|
||||
executedBatches []*LLMBatch
|
||||
pendingBatches []*LLMBatch
|
||||
stats *LLMBatchStats
|
||||
}
|
||||
|
||||
// LLMBatchStats tracks LLM batching statistics
|
||||
type LLMBatchStats struct {
|
||||
TotalRequests int
|
||||
TotalBatches int
|
||||
AvgRequestsPerBatch float64
|
||||
APICallsSaved int // Total API calls saved (individual requests - batches)
|
||||
TotalTokens int
|
||||
TotalExecutionTime time.Duration
|
||||
}
|
||||
|
||||
// NewLLMBatcher creates a new LLM batcher
|
||||
func NewLLMBatcher(maxBatchSize int, maxBatchAge time.Duration) *LLMBatcher {
|
||||
if maxBatchSize <= 0 {
|
||||
maxBatchSize = 10
|
||||
}
|
||||
if maxBatchAge <= 0 {
|
||||
maxBatchAge = 2 * time.Second
|
||||
}
|
||||
|
||||
return &LLMBatcher{
|
||||
queue: make([]*LLMRequest, 0),
|
||||
maxBatchSize: maxBatchSize,
|
||||
maxBatchAge: maxBatchAge,
|
||||
lastFlushTime: time.Now(),
|
||||
executedBatches: make([]*LLMBatch, 0),
|
||||
pendingBatches: make([]*LLMBatch, 0),
|
||||
stats: &LLMBatchStats{
|
||||
TotalRequests: 0,
|
||||
TotalBatches: 0,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// Enqueue adds an LLM request to the queue
|
||||
func (lb *LLMBatcher) Enqueue(req *LLMRequest) {
|
||||
if req == nil {
|
||||
return
|
||||
}
|
||||
|
||||
if req.ID == "" {
|
||||
req.ID = fmt.Sprintf("req-%d", time.Now().UnixNano())
|
||||
}
|
||||
|
||||
req.Timestamp = time.Now()
|
||||
if req.ResultCh == nil {
|
||||
req.ResultCh = make(chan *LLMResult, 1)
|
||||
}
|
||||
|
||||
lb.mu.Lock()
|
||||
defer lb.mu.Unlock()
|
||||
|
||||
lb.queue = append(lb.queue, req)
|
||||
lb.stats.TotalRequests++
|
||||
|
||||
// Auto-flush if batch is full
|
||||
if len(lb.queue) >= lb.maxBatchSize {
|
||||
lb.flushLocked()
|
||||
}
|
||||
}
|
||||
|
||||
// flushLocked creates a batch from queued requests (must be called with lock held)
|
||||
func (lb *LLMBatcher) flushLocked() {
|
||||
if len(lb.queue) == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
// Group by type and model
|
||||
groups := make(map[string][]*LLMRequest)
|
||||
for _, req := range lb.queue {
|
||||
key := fmt.Sprintf("%s:%s", req.Type, req.Model)
|
||||
groups[key] = append(groups[key], req)
|
||||
}
|
||||
|
||||
// Create batch for each group
|
||||
for key, reqs := range groups {
|
||||
batch := &LLMBatch{
|
||||
ID: fmt.Sprintf("batch-%d", lb.stats.TotalBatches),
|
||||
Requests: reqs,
|
||||
Model: reqs[0].Model,
|
||||
Type: reqs[0].Type,
|
||||
CreatedAt: time.Now(),
|
||||
Status: "pending",
|
||||
Results: make(map[string]*LLMResult),
|
||||
}
|
||||
|
||||
lb.pendingBatches = append(lb.pendingBatches, batch)
|
||||
lb.stats.TotalBatches++
|
||||
_ = key // Silence unused variable warning
|
||||
}
|
||||
|
||||
lb.queue = make([]*LLMRequest, 0)
|
||||
lb.lastFlushTime = time.Now()
|
||||
}
|
||||
|
||||
// Flush manually flushes the current queue
|
||||
func (lb *LLMBatcher) Flush() {
|
||||
lb.mu.Lock()
|
||||
defer lb.mu.Unlock()
|
||||
|
||||
lb.flushLocked()
|
||||
}
|
||||
|
||||
// GetPendingBatch returns the next pending batch without removing it
|
||||
func (lb *LLMBatcher) GetPendingBatch() *LLMBatch {
|
||||
lb.mu.RLock()
|
||||
defer lb.mu.RUnlock()
|
||||
|
||||
if len(lb.pendingBatches) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
return lb.pendingBatches[0]
|
||||
}
|
||||
|
||||
// MarkBatchExecuting marks a batch as executing
|
||||
func (lb *LLMBatcher) MarkBatchExecuting(batchID string) {
|
||||
lb.mu.Lock()
|
||||
defer lb.mu.Unlock()
|
||||
|
||||
for _, batch := range lb.pendingBatches {
|
||||
if batch.ID == batchID {
|
||||
batch.Status = "executing"
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// MarkBatchCompleted marks a batch as completed and delivers results
|
||||
func (lb *LLMBatcher) MarkBatchCompleted(batchID string, results map[string]*LLMResult) {
|
||||
lb.mu.Lock()
|
||||
defer lb.mu.Unlock()
|
||||
|
||||
var idx int
|
||||
var found *LLMBatch
|
||||
for i, batch := range lb.pendingBatches {
|
||||
if batch.ID == batchID {
|
||||
idx = i
|
||||
found = batch
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if found != nil {
|
||||
found.Status = "completed"
|
||||
found.ExecutedAt = time.Now()
|
||||
found.ExecutionTime = found.ExecutedAt.Sub(found.CreatedAt)
|
||||
found.Results = results
|
||||
|
||||
// Deliver results to request channels
|
||||
for _, req := range found.Requests {
|
||||
if result, exists := results[req.ID]; exists {
|
||||
select {
|
||||
case req.ResultCh <- result:
|
||||
default:
|
||||
// Channel not ready or closed
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Update stats
|
||||
lb.stats.TotalTokens += countTokensInBatch(found)
|
||||
lb.stats.TotalExecutionTime += found.ExecutionTime
|
||||
|
||||
// Move to executed batches
|
||||
lb.executedBatches = append(lb.executedBatches, found)
|
||||
lb.pendingBatches = append(lb.pendingBatches[:idx], lb.pendingBatches[idx+1:]...)
|
||||
}
|
||||
}
|
||||
|
||||
// MarkBatchFailed marks a batch as failed
|
||||
func (lb *LLMBatcher) MarkBatchFailed(batchID string, err error) {
|
||||
lb.mu.Lock()
|
||||
defer lb.mu.Unlock()
|
||||
|
||||
var found *LLMBatch
|
||||
for _, batch := range lb.pendingBatches {
|
||||
if batch.ID == batchID {
|
||||
found = batch
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if found != nil {
|
||||
found.Status = "failed"
|
||||
found.Error = err
|
||||
found.ExecutedAt = time.Now()
|
||||
|
||||
// Deliver errors to request channels
|
||||
for _, req := range found.Requests {
|
||||
result := &LLMResult{
|
||||
RequestID: req.ID,
|
||||
Error: err,
|
||||
}
|
||||
|
||||
select {
|
||||
case req.ResultCh <- result:
|
||||
default:
|
||||
// Channel not ready or closed
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// GetStats returns batching statistics
|
||||
func (lb *LLMBatcher) GetStats() *LLMBatchStats {
|
||||
lb.mu.RLock()
|
||||
defer lb.mu.RUnlock()
|
||||
|
||||
stats := *lb.stats
|
||||
if stats.TotalBatches > 0 {
|
||||
stats.AvgRequestsPerBatch = float64(stats.TotalRequests) / float64(stats.TotalBatches)
|
||||
// API calls saved: total requests - total batches
|
||||
stats.APICallsSaved = stats.TotalRequests - stats.TotalBatches
|
||||
}
|
||||
|
||||
return &stats
|
||||
}
|
||||
|
||||
// QueueSize returns current queue size
|
||||
func (lb *LLMBatcher) QueueSize() int {
|
||||
lb.mu.RLock()
|
||||
defer lb.mu.RUnlock()
|
||||
|
||||
return len(lb.queue)
|
||||
}
|
||||
|
||||
// PendingBatchCount returns number of pending batches
|
||||
func (lb *LLMBatcher) PendingBatchCount() int {
|
||||
lb.mu.RLock()
|
||||
defer lb.mu.RUnlock()
|
||||
|
||||
return len(lb.pendingBatches)
|
||||
}
|
||||
|
||||
// GetBatchByID returns a batch by ID
|
||||
func (lb *LLMBatcher) GetBatchByID(batchID string) *LLMBatch {
|
||||
lb.mu.RLock()
|
||||
defer lb.mu.RUnlock()
|
||||
|
||||
for _, batch := range lb.pendingBatches {
|
||||
if batch.ID == batchID {
|
||||
return batch
|
||||
}
|
||||
}
|
||||
|
||||
for _, batch := range lb.executedBatches {
|
||||
if batch.ID == batchID {
|
||||
return batch
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// TimeSinceLastFlush returns time since last flush
|
||||
func (lb *LLMBatcher) TimeSinceLastFlush() time.Duration {
|
||||
lb.mu.RLock()
|
||||
defer lb.mu.RUnlock()
|
||||
|
||||
return time.Since(lb.lastFlushTime)
|
||||
}
|
||||
|
||||
// ShouldFlush checks if queue should be flushed based on age
|
||||
func (lb *LLMBatcher) ShouldFlush() bool {
|
||||
lb.mu.RLock()
|
||||
defer lb.mu.RUnlock()
|
||||
|
||||
if len(lb.queue) == 0 {
|
||||
return false
|
||||
}
|
||||
|
||||
return time.Since(lb.lastFlushTime) >= lb.maxBatchAge
|
||||
}
|
||||
|
||||
// GetExecutedBatches returns all executed batches
|
||||
func (lb *LLMBatcher) GetExecutedBatches() []*LLMBatch {
|
||||
lb.mu.RLock()
|
||||
defer lb.mu.RUnlock()
|
||||
|
||||
result := make([]*LLMBatch, len(lb.executedBatches))
|
||||
copy(result, lb.executedBatches)
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
// Clear clears all pending operations
|
||||
func (lb *LLMBatcher) Clear() {
|
||||
lb.mu.Lock()
|
||||
defer lb.mu.Unlock()
|
||||
|
||||
lb.queue = make([]*LLMRequest, 0)
|
||||
lb.pendingBatches = make([]*LLMBatch, 0)
|
||||
lb.executedBatches = make([]*LLMBatch, 0)
|
||||
}
|
||||
|
||||
// countTokensInBatch counts total tokens in a batch
|
||||
func countTokensInBatch(batch *LLMBatch) int {
|
||||
total := 0
|
||||
for _, result := range batch.Results {
|
||||
total += result.TokenCount
|
||||
}
|
||||
return total
|
||||
}
|
||||
|
||||
// GetBatchInfo returns human-readable batch information
|
||||
func (batch *LLMBatch) GetInfo() map[string]interface{} {
|
||||
return map[string]interface{}{
|
||||
"id": batch.ID,
|
||||
"type": batch.Type,
|
||||
"model": batch.Model,
|
||||
"status": batch.Status,
|
||||
"request_count": len(batch.Requests),
|
||||
"created_at": batch.CreatedAt,
|
||||
"executed_at": batch.ExecutedAt,
|
||||
"duration": batch.ExecutionTime,
|
||||
"error": batch.Error,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,436 @@
|
||||
package batching
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
func TestLLMNewBatcher(t *testing.T) {
|
||||
batcher := NewLLMBatcher(10, 5*time.Second)
|
||||
assert.NotNil(t, batcher)
|
||||
assert.Equal(t, 0, batcher.QueueSize())
|
||||
}
|
||||
|
||||
func TestLLMEnqueueRequest(t *testing.T) {
|
||||
batcher := NewLLMBatcher(10, 5*time.Second)
|
||||
|
||||
req := &LLMRequest{
|
||||
ID: "req-1",
|
||||
Type: "implementer",
|
||||
Model: "claude-opus",
|
||||
Prompt: "Generate code",
|
||||
}
|
||||
|
||||
batcher.Enqueue(req)
|
||||
assert.Equal(t, 1, batcher.QueueSize())
|
||||
}
|
||||
|
||||
func TestLLMEnqueueMultipleRequests(t *testing.T) {
|
||||
batcher := NewLLMBatcher(10, 5*time.Second)
|
||||
|
||||
for i := 0; i < 5; i++ {
|
||||
req := &LLMRequest{
|
||||
Type: "implementer",
|
||||
Model: "claude-opus",
|
||||
Prompt: "Prompt",
|
||||
}
|
||||
batcher.Enqueue(req)
|
||||
}
|
||||
|
||||
assert.Equal(t, 5, batcher.QueueSize())
|
||||
}
|
||||
|
||||
func TestLLMAutoFlushOnMaxBatchSize(t *testing.T) {
|
||||
batcher := NewLLMBatcher(5, 10*time.Second)
|
||||
|
||||
for i := 0; i < 5; i++ {
|
||||
req := &LLMRequest{
|
||||
Type: "implementer",
|
||||
Model: "claude-opus",
|
||||
Prompt: "Prompt",
|
||||
}
|
||||
batcher.Enqueue(req)
|
||||
}
|
||||
|
||||
assert.Equal(t, 0, batcher.QueueSize())
|
||||
assert.Equal(t, 1, batcher.PendingBatchCount())
|
||||
}
|
||||
|
||||
func TestLLMManualFlush(t *testing.T) {
|
||||
batcher := NewLLMBatcher(10, 5*time.Second)
|
||||
|
||||
req := &LLMRequest{
|
||||
Type: "implementer",
|
||||
Model: "claude-opus",
|
||||
Prompt: "Prompt",
|
||||
}
|
||||
batcher.Enqueue(req)
|
||||
|
||||
batcher.Flush()
|
||||
assert.Equal(t, 0, batcher.QueueSize())
|
||||
assert.Equal(t, 1, batcher.PendingBatchCount())
|
||||
}
|
||||
|
||||
func TestLLMGetPendingBatch(t *testing.T) {
|
||||
batcher := NewLLMBatcher(10, 5*time.Second)
|
||||
|
||||
req := &LLMRequest{
|
||||
Type: "implementer",
|
||||
Model: "claude-opus",
|
||||
Prompt: "Prompt",
|
||||
}
|
||||
batcher.Enqueue(req)
|
||||
batcher.Flush()
|
||||
|
||||
batch := batcher.GetPendingBatch()
|
||||
assert.NotNil(t, batch)
|
||||
assert.Equal(t, 1, len(batch.Requests))
|
||||
assert.Equal(t, "implementer", batch.Type)
|
||||
}
|
||||
|
||||
func TestLLMMarkBatchExecuting(t *testing.T) {
|
||||
batcher := NewLLMBatcher(10, 5*time.Second)
|
||||
|
||||
req := &LLMRequest{Type: "implementer", Model: "claude-opus", Prompt: "test"}
|
||||
batcher.Enqueue(req)
|
||||
batcher.Flush()
|
||||
|
||||
batch := batcher.GetPendingBatch()
|
||||
batcher.MarkBatchExecuting(batch.ID)
|
||||
|
||||
updated := batcher.GetBatchByID(batch.ID)
|
||||
assert.Equal(t, "executing", updated.Status)
|
||||
}
|
||||
|
||||
func TestLLMMarkBatchCompleted(t *testing.T) {
|
||||
batcher := NewLLMBatcher(10, 5*time.Second)
|
||||
|
||||
req := &LLMRequest{
|
||||
ID: "req-1",
|
||||
Type: "implementer",
|
||||
Model: "claude-opus",
|
||||
Prompt: "test",
|
||||
}
|
||||
batcher.Enqueue(req)
|
||||
batcher.Flush()
|
||||
|
||||
batch := batcher.GetPendingBatch()
|
||||
|
||||
results := map[string]*LLMResult{
|
||||
"req-1": {
|
||||
RequestID: "req-1",
|
||||
Response: "Generated code",
|
||||
TokenCount: 100,
|
||||
},
|
||||
}
|
||||
|
||||
batcher.MarkBatchCompleted(batch.ID, results)
|
||||
|
||||
executed := batcher.GetExecutedBatches()
|
||||
assert.Equal(t, 1, len(executed))
|
||||
assert.Equal(t, "completed", executed[0].Status)
|
||||
}
|
||||
|
||||
func TestLLMMarkBatchFailed(t *testing.T) {
|
||||
batcher := NewLLMBatcher(10, 5*time.Second)
|
||||
|
||||
req := &LLMRequest{
|
||||
ID: "req-1",
|
||||
Type: "implementer",
|
||||
Model: "claude-opus",
|
||||
}
|
||||
batcher.Enqueue(req)
|
||||
batcher.Flush()
|
||||
|
||||
batch := batcher.GetPendingBatch()
|
||||
testErr := assert.AnError
|
||||
batcher.MarkBatchFailed(batch.ID, testErr)
|
||||
|
||||
failed := batcher.GetBatchByID(batch.ID)
|
||||
assert.Equal(t, "failed", failed.Status)
|
||||
assert.Error(t, failed.Error)
|
||||
}
|
||||
|
||||
func TestLLMGroupByTypeAndModel(t *testing.T) {
|
||||
batcher := NewLLMBatcher(100, 5*time.Second)
|
||||
|
||||
// Add requests of different types
|
||||
for i := 0; i < 3; i++ {
|
||||
req := &LLMRequest{
|
||||
Type: "implementer",
|
||||
Model: "claude-opus",
|
||||
Prompt: "test",
|
||||
}
|
||||
batcher.Enqueue(req)
|
||||
}
|
||||
|
||||
for i := 0; i < 2; i++ {
|
||||
req := &LLMRequest{
|
||||
Type: "judge",
|
||||
Model: "claude-opus",
|
||||
Prompt: "test",
|
||||
}
|
||||
batcher.Enqueue(req)
|
||||
}
|
||||
|
||||
batcher.Flush()
|
||||
|
||||
// Should create 2 batches (one for implementer, one for judge)
|
||||
assert.Equal(t, 2, batcher.PendingBatchCount())
|
||||
}
|
||||
|
||||
func TestLLMGetStats(t *testing.T) {
|
||||
batcher := NewLLMBatcher(5, 5*time.Second)
|
||||
|
||||
// Add 10 requests (will create 2 batches)
|
||||
for i := 0; i < 10; i++ {
|
||||
req := &LLMRequest{
|
||||
Type: "implementer",
|
||||
Model: "claude-opus",
|
||||
Prompt: "test",
|
||||
}
|
||||
batcher.Enqueue(req)
|
||||
}
|
||||
|
||||
stats := batcher.GetStats()
|
||||
assert.Equal(t, 10, stats.TotalRequests)
|
||||
assert.Equal(t, 2, stats.TotalBatches)
|
||||
assert.Equal(t, 5.0, stats.AvgRequestsPerBatch)
|
||||
// 10 requests in 2 batches saves 8 API calls
|
||||
assert.Equal(t, 8, stats.APICallsSaved)
|
||||
}
|
||||
|
||||
func TestLLMResultDelivery(t *testing.T) {
|
||||
batcher := NewLLMBatcher(10, 5*time.Second)
|
||||
|
||||
req := &LLMRequest{
|
||||
ID: "req-1",
|
||||
Type: "implementer",
|
||||
Model: "claude-opus",
|
||||
Prompt: "test",
|
||||
ResultCh: make(chan *LLMResult, 1),
|
||||
}
|
||||
|
||||
batcher.Enqueue(req)
|
||||
batcher.Flush()
|
||||
|
||||
batch := batcher.GetPendingBatch()
|
||||
|
||||
results := map[string]*LLMResult{
|
||||
"req-1": {
|
||||
RequestID: "req-1",
|
||||
Response: "Response",
|
||||
TokenCount: 50,
|
||||
},
|
||||
}
|
||||
|
||||
batcher.MarkBatchCompleted(batch.ID, results)
|
||||
|
||||
// Check if result was delivered to channel
|
||||
select {
|
||||
case result := <-req.ResultCh:
|
||||
assert.NotNil(t, result)
|
||||
assert.Equal(t, "Response", result.Response)
|
||||
case <-time.After(1 * time.Second):
|
||||
t.Fatal("Result not delivered")
|
||||
}
|
||||
}
|
||||
|
||||
func TestLLMMultipleBatches(t *testing.T) {
|
||||
batcher := NewLLMBatcher(3, 5*time.Second)
|
||||
|
||||
// Create 3 batches (3 requests each)
|
||||
for batch := 0; batch < 3; batch++ {
|
||||
for i := 0; i < 3; i++ {
|
||||
req := &LLMRequest{
|
||||
Type: "implementer",
|
||||
Model: "claude-opus",
|
||||
Prompt: "test",
|
||||
}
|
||||
batcher.Enqueue(req)
|
||||
}
|
||||
}
|
||||
|
||||
assert.Equal(t, 3, batcher.PendingBatchCount())
|
||||
}
|
||||
|
||||
func TestLLMQueueSize(t *testing.T) {
|
||||
batcher := NewLLMBatcher(10, 5*time.Second)
|
||||
|
||||
req := &LLMRequest{Type: "implementer", Model: "claude-opus"}
|
||||
batcher.Enqueue(req)
|
||||
|
||||
assert.Equal(t, 1, batcher.QueueSize())
|
||||
}
|
||||
|
||||
func TestLLMPendingBatchCount(t *testing.T) {
|
||||
batcher := NewLLMBatcher(10, 5*time.Second)
|
||||
|
||||
req := &LLMRequest{Type: "implementer", Model: "claude-opus"}
|
||||
batcher.Enqueue(req)
|
||||
batcher.Flush()
|
||||
|
||||
assert.Equal(t, 1, batcher.PendingBatchCount())
|
||||
}
|
||||
|
||||
func TestLLMGetExecutedBatches(t *testing.T) {
|
||||
batcher := NewLLMBatcher(10, 5*time.Second)
|
||||
|
||||
for i := 0; i < 2; i++ {
|
||||
req := &LLMRequest{Type: "implementer", Model: "claude-opus"}
|
||||
batcher.Enqueue(req)
|
||||
batcher.Flush()
|
||||
|
||||
batch := batcher.GetPendingBatch()
|
||||
batcher.MarkBatchCompleted(batch.ID, make(map[string]*LLMResult))
|
||||
}
|
||||
|
||||
executed := batcher.GetExecutedBatches()
|
||||
assert.Equal(t, 2, len(executed))
|
||||
}
|
||||
|
||||
func TestLLMTimeSinceLastFlush(t *testing.T) {
|
||||
batcher := NewLLMBatcher(10, 5*time.Second)
|
||||
|
||||
req := &LLMRequest{Type: "implementer", Model: "claude-opus"}
|
||||
batcher.Enqueue(req)
|
||||
batcher.Flush()
|
||||
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
elapsed := batcher.TimeSinceLastFlush()
|
||||
|
||||
assert.Greater(t, elapsed, 50*time.Millisecond)
|
||||
assert.Less(t, elapsed, 200*time.Millisecond)
|
||||
}
|
||||
|
||||
func TestLLMShouldFlush(t *testing.T) {
|
||||
batcher := NewLLMBatcher(100, 100*time.Millisecond)
|
||||
|
||||
req := &LLMRequest{Type: "implementer", Model: "claude-opus"}
|
||||
batcher.Enqueue(req)
|
||||
|
||||
// Should not flush yet
|
||||
assert.False(t, batcher.ShouldFlush())
|
||||
|
||||
// Wait for age to exceed
|
||||
time.Sleep(150 * time.Millisecond)
|
||||
assert.True(t, batcher.ShouldFlush())
|
||||
}
|
||||
|
||||
func TestLLMClear(t *testing.T) {
|
||||
batcher := NewLLMBatcher(10, 5*time.Second)
|
||||
|
||||
req := &LLMRequest{Type: "implementer", Model: "claude-opus"}
|
||||
batcher.Enqueue(req)
|
||||
batcher.Flush()
|
||||
|
||||
batcher.Clear()
|
||||
assert.Equal(t, 0, batcher.QueueSize())
|
||||
assert.Equal(t, 0, batcher.PendingBatchCount())
|
||||
}
|
||||
|
||||
func TestLLMGetBatchByID(t *testing.T) {
|
||||
batcher := NewLLMBatcher(10, 5*time.Second)
|
||||
|
||||
req := &LLMRequest{Type: "implementer", Model: "claude-opus"}
|
||||
batcher.Enqueue(req)
|
||||
batcher.Flush()
|
||||
|
||||
batch := batcher.GetPendingBatch()
|
||||
retrieved := batcher.GetBatchByID(batch.ID)
|
||||
|
||||
assert.NotNil(t, retrieved)
|
||||
assert.Equal(t, batch.ID, retrieved.ID)
|
||||
}
|
||||
|
||||
func TestLLMGetBatchInfo(t *testing.T) {
|
||||
batch := &LLMBatch{
|
||||
ID: "batch-1",
|
||||
Type: "implementer",
|
||||
Model: "claude-opus",
|
||||
Status: "completed",
|
||||
CreatedAt: time.Now(),
|
||||
}
|
||||
|
||||
info := batch.GetInfo()
|
||||
assert.Equal(t, "batch-1", info["id"])
|
||||
assert.Equal(t, "implementer", info["type"])
|
||||
assert.Equal(t, "claude-opus", info["model"])
|
||||
assert.Equal(t, "completed", info["status"])
|
||||
}
|
||||
|
||||
func TestLLMEnqueueNil(t *testing.T) {
|
||||
batcher := NewLLMBatcher(10, 5*time.Second)
|
||||
|
||||
batcher.Enqueue(nil)
|
||||
assert.Equal(t, 0, batcher.QueueSize())
|
||||
}
|
||||
|
||||
func TestLLMAutoIDGeneration(t *testing.T) {
|
||||
req := &LLMRequest{
|
||||
Type: "implementer",
|
||||
Model: "claude-opus",
|
||||
Prompt: "test",
|
||||
}
|
||||
|
||||
batcher := NewLLMBatcher(10, 5*time.Second)
|
||||
batcher.Enqueue(req)
|
||||
|
||||
assert.NotEmpty(t, req.ID)
|
||||
}
|
||||
|
||||
func TestLLMTokenCounting(t *testing.T) {
|
||||
batcher := NewLLMBatcher(10, 5*time.Second)
|
||||
|
||||
req := &LLMRequest{
|
||||
ID: "req-1",
|
||||
Type: "implementer",
|
||||
Model: "claude-opus",
|
||||
Prompt: "test",
|
||||
}
|
||||
batcher.Enqueue(req)
|
||||
batcher.Flush()
|
||||
|
||||
batch := batcher.GetPendingBatch()
|
||||
|
||||
results := map[string]*LLMResult{
|
||||
"req-1": {
|
||||
RequestID: "req-1",
|
||||
Response: "Response",
|
||||
TokenCount: 500,
|
||||
},
|
||||
}
|
||||
|
||||
batcher.MarkBatchCompleted(batch.ID, results)
|
||||
|
||||
stats := batcher.GetStats()
|
||||
assert.Equal(t, 500, stats.TotalTokens)
|
||||
}
|
||||
|
||||
func BenchmarkLLMEnqueue(b *testing.B) {
|
||||
batcher := NewLLMBatcher(1000, 10*time.Second)
|
||||
|
||||
for i := 0; i < b.N; i++ {
|
||||
req := &LLMRequest{
|
||||
Type: "implementer",
|
||||
Model: "claude-opus",
|
||||
Prompt: "test",
|
||||
}
|
||||
batcher.Enqueue(req)
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkLLMFlush(b *testing.B) {
|
||||
batcher := NewLLMBatcher(1000, 10*time.Second)
|
||||
|
||||
for i := 0; i < b.N; i++ {
|
||||
req := &LLMRequest{Type: "implementer", Model: "claude-opus"}
|
||||
batcher.Enqueue(req)
|
||||
|
||||
if (i + 1) % 100 == 0 {
|
||||
batcher.Flush()
|
||||
}
|
||||
}
|
||||
}
|
||||
Vendored
+319
@@ -0,0 +1,319 @@
|
||||
package cache
|
||||
|
||||
import (
|
||||
"crypto/md5"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// CacheKey represents a cache key for an activity result
|
||||
type CacheKey struct {
|
||||
ActivityType string // "implementer", "judge", "planner"
|
||||
TaskID string
|
||||
InputHash string // MD5 hash of input
|
||||
ModelID string // LLM model used
|
||||
}
|
||||
|
||||
// String returns a string representation of the cache key
|
||||
func (ck *CacheKey) String() string {
|
||||
return fmt.Sprintf("%s:%s:%s:%s", ck.ActivityType, ck.TaskID, ck.InputHash, ck.ModelID)
|
||||
}
|
||||
|
||||
// CacheEntry represents a cached activity result
|
||||
type CacheEntry struct {
|
||||
Key CacheKey `json:"key"`
|
||||
Result map[string]interface{} `json:"result"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
HitCount int `json:"hit_count"`
|
||||
Metadata map[string]interface{} `json:"metadata,omitempty"`
|
||||
}
|
||||
|
||||
// ResultCache caches activity results to avoid redundant computations
|
||||
type ResultCache struct {
|
||||
mu sync.RWMutex
|
||||
basePath string
|
||||
cache map[string]*CacheEntry
|
||||
maxSize int
|
||||
ttl time.Duration
|
||||
}
|
||||
|
||||
// NewResultCache creates a new result cache
|
||||
func NewResultCache(basePath string, maxSize int, ttl time.Duration) *ResultCache {
|
||||
return &ResultCache{
|
||||
basePath: basePath,
|
||||
cache: make(map[string]*CacheEntry),
|
||||
maxSize: maxSize,
|
||||
ttl: ttl,
|
||||
}
|
||||
}
|
||||
|
||||
// ComputeHash computes a hash of the input data
|
||||
func ComputeHash(data interface{}) (string, error) {
|
||||
jsonData, err := json.Marshal(data)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
hash := md5.Sum(jsonData)
|
||||
return fmt.Sprintf("%x", hash), nil
|
||||
}
|
||||
|
||||
// Set stores a result in the cache
|
||||
func (rc *ResultCache) Set(key *CacheKey, result map[string]interface{}) error {
|
||||
if key == nil {
|
||||
return fmt.Errorf("cache key cannot be nil")
|
||||
}
|
||||
|
||||
rc.mu.Lock()
|
||||
defer rc.mu.Unlock()
|
||||
|
||||
keyStr := key.String()
|
||||
|
||||
entry := &CacheEntry{
|
||||
Key: *key,
|
||||
Result: result,
|
||||
CreatedAt: time.Now(),
|
||||
Metadata: make(map[string]interface{}),
|
||||
}
|
||||
|
||||
// Check size limit
|
||||
if len(rc.cache) >= rc.maxSize && rc.cache[keyStr] == nil {
|
||||
// Evict oldest entry (simple FIFO)
|
||||
var oldestKey string
|
||||
var oldestTime time.Time
|
||||
|
||||
for k, v := range rc.cache {
|
||||
if oldestTime.IsZero() || v.CreatedAt.Before(oldestTime) {
|
||||
oldestKey = k
|
||||
oldestTime = v.CreatedAt
|
||||
}
|
||||
}
|
||||
|
||||
if oldestKey != "" {
|
||||
delete(rc.cache, oldestKey)
|
||||
}
|
||||
}
|
||||
|
||||
rc.cache[keyStr] = entry
|
||||
return rc.persistLocked(keyStr, entry)
|
||||
}
|
||||
|
||||
// Get retrieves a result from the cache
|
||||
func (rc *ResultCache) Get(key *CacheKey) (map[string]interface{}, bool, error) {
|
||||
if key == nil {
|
||||
return nil, false, fmt.Errorf("cache key cannot be nil")
|
||||
}
|
||||
|
||||
rc.mu.Lock()
|
||||
defer rc.mu.Unlock()
|
||||
|
||||
keyStr := key.String()
|
||||
entry, exists := rc.cache[keyStr]
|
||||
|
||||
if !exists {
|
||||
return nil, false, nil
|
||||
}
|
||||
|
||||
// Check TTL
|
||||
if rc.ttl > 0 && time.Since(entry.CreatedAt) > rc.ttl {
|
||||
delete(rc.cache, keyStr)
|
||||
return nil, false, nil
|
||||
}
|
||||
|
||||
// Increment hit count
|
||||
entry.HitCount++
|
||||
_ = rc.persistLocked(keyStr, entry)
|
||||
|
||||
return entry.Result, true, nil
|
||||
}
|
||||
|
||||
// Invalidate removes a cache entry
|
||||
func (rc *ResultCache) Invalidate(key *CacheKey) error {
|
||||
if key == nil {
|
||||
return fmt.Errorf("cache key cannot be nil")
|
||||
}
|
||||
|
||||
rc.mu.Lock()
|
||||
defer rc.mu.Unlock()
|
||||
|
||||
keyStr := key.String()
|
||||
delete(rc.cache, keyStr)
|
||||
|
||||
// Delete from disk
|
||||
cacheFile := filepath.Join(rc.basePath, "cache", fmt.Sprintf("%s.json", keyStr))
|
||||
_ = os.Remove(cacheFile)
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// Clear clears all cache entries
|
||||
func (rc *ResultCache) Clear() error {
|
||||
rc.mu.Lock()
|
||||
defer rc.mu.Unlock()
|
||||
|
||||
rc.cache = make(map[string]*CacheEntry)
|
||||
|
||||
// Clear disk cache
|
||||
cacheDir := filepath.Join(rc.basePath, "cache")
|
||||
_ = os.RemoveAll(cacheDir)
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetStats returns cache statistics
|
||||
func (rc *ResultCache) GetStats() map[string]interface{} {
|
||||
rc.mu.RLock()
|
||||
defer rc.mu.RUnlock()
|
||||
|
||||
totalHits := 0
|
||||
for _, entry := range rc.cache {
|
||||
totalHits += entry.HitCount
|
||||
}
|
||||
|
||||
return map[string]interface{}{
|
||||
"size": len(rc.cache),
|
||||
"max_size": rc.maxSize,
|
||||
"total_hits": totalHits,
|
||||
"usage_ratio": float64(len(rc.cache)) / float64(rc.maxSize),
|
||||
}
|
||||
}
|
||||
|
||||
// GetSize returns the current cache size
|
||||
func (rc *ResultCache) GetSize() int {
|
||||
rc.mu.RLock()
|
||||
defer rc.mu.RUnlock()
|
||||
|
||||
return len(rc.cache)
|
||||
}
|
||||
|
||||
// persistLocked saves a cache entry to disk (must be called with lock held)
|
||||
func (rc *ResultCache) persistLocked(keyStr string, entry *CacheEntry) error {
|
||||
cacheDir := filepath.Join(rc.basePath, "cache")
|
||||
|
||||
// Create directory if it doesn't exist
|
||||
if err := os.MkdirAll(cacheDir, 0755); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
cacheFile := filepath.Join(cacheDir, fmt.Sprintf("%s.json", keyStr))
|
||||
|
||||
data, err := json.MarshalIndent(entry, "", " ")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return os.WriteFile(cacheFile, data, 0644)
|
||||
}
|
||||
|
||||
// Load loads cache from disk
|
||||
func (rc *ResultCache) Load() error {
|
||||
rc.mu.Lock()
|
||||
defer rc.mu.Unlock()
|
||||
|
||||
cacheDir := filepath.Join(rc.basePath, "cache")
|
||||
entries, err := os.ReadDir(cacheDir)
|
||||
if err != nil {
|
||||
if os.IsNotExist(err) {
|
||||
return nil // Cache doesn't exist yet
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
for _, entry := range entries {
|
||||
if entry.IsDir() {
|
||||
continue
|
||||
}
|
||||
|
||||
filePath := filepath.Join(cacheDir, entry.Name())
|
||||
data, err := os.ReadFile(filePath)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
|
||||
var cacheEntry CacheEntry
|
||||
if err := json.Unmarshal(data, &cacheEntry); err != nil {
|
||||
continue
|
||||
}
|
||||
|
||||
// Skip expired entries
|
||||
if rc.ttl > 0 && time.Since(cacheEntry.CreatedAt) > rc.ttl {
|
||||
continue
|
||||
}
|
||||
|
||||
keyStr := cacheEntry.Key.String()
|
||||
rc.cache[keyStr] = &cacheEntry
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// InvalidateByActivity invalidates all cache entries for an activity type
|
||||
func (rc *ResultCache) InvalidateByActivity(activityType string) error {
|
||||
rc.mu.Lock()
|
||||
defer rc.mu.Unlock()
|
||||
|
||||
keysToDelete := make([]string, 0)
|
||||
for keyStr, entry := range rc.cache {
|
||||
if entry.Key.ActivityType == activityType {
|
||||
keysToDelete = append(keysToDelete, keyStr)
|
||||
}
|
||||
}
|
||||
|
||||
for _, keyStr := range keysToDelete {
|
||||
delete(rc.cache, keyStr)
|
||||
|
||||
// Delete from disk
|
||||
cacheFile := filepath.Join(rc.basePath, "cache", fmt.Sprintf("%s.json", keyStr))
|
||||
_ = os.Remove(cacheFile)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// InvalidateByTask invalidates all cache entries for a task
|
||||
func (rc *ResultCache) InvalidateByTask(taskID string) error {
|
||||
rc.mu.Lock()
|
||||
defer rc.mu.Unlock()
|
||||
|
||||
keysToDelete := make([]string, 0)
|
||||
for keyStr, entry := range rc.cache {
|
||||
if entry.Key.TaskID == taskID {
|
||||
keysToDelete = append(keysToDelete, keyStr)
|
||||
}
|
||||
}
|
||||
|
||||
for _, keyStr := range keysToDelete {
|
||||
delete(rc.cache, keyStr)
|
||||
|
||||
// Delete from disk
|
||||
cacheFile := filepath.Join(rc.basePath, "cache", fmt.Sprintf("%s.json", keyStr))
|
||||
_ = os.Remove(cacheFile)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetHitRate returns the cache hit rate
|
||||
func (rc *ResultCache) GetHitRate() (float64, int) {
|
||||
rc.mu.RLock()
|
||||
defer rc.mu.RUnlock()
|
||||
|
||||
if len(rc.cache) == 0 {
|
||||
return 0, 0
|
||||
}
|
||||
|
||||
totalHits := 0
|
||||
for _, entry := range rc.cache {
|
||||
totalHits += entry.HitCount
|
||||
}
|
||||
|
||||
if totalHits == 0 {
|
||||
return 0, len(rc.cache)
|
||||
}
|
||||
|
||||
return float64(totalHits) / float64(len(rc.cache)), len(rc.cache)
|
||||
}
|
||||
Vendored
+316
@@ -0,0 +1,316 @@
|
||||
package cache
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
func TestCacheKeyString(t *testing.T) {
|
||||
key := &CacheKey{
|
||||
ActivityType: "implementer",
|
||||
TaskID: "T1.1",
|
||||
InputHash: "abc123",
|
||||
ModelID: "claude-opus",
|
||||
}
|
||||
|
||||
keyStr := key.String()
|
||||
assert.Contains(t, keyStr, "implementer")
|
||||
assert.Contains(t, keyStr, "T1.1")
|
||||
assert.Contains(t, keyStr, "abc123")
|
||||
assert.Contains(t, keyStr, "claude-opus")
|
||||
}
|
||||
|
||||
func TestComputeHash(t *testing.T) {
|
||||
data := map[string]interface{}{
|
||||
"task": "T1.1",
|
||||
"code": "package main",
|
||||
}
|
||||
|
||||
hash1, err := ComputeHash(data)
|
||||
assert.NoError(t, err)
|
||||
assert.NotEmpty(t, hash1)
|
||||
|
||||
hash2, err := ComputeHash(data)
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, hash1, hash2)
|
||||
}
|
||||
|
||||
func TestSetAndGet(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
cache := NewResultCache(tmpDir, 100, 0)
|
||||
|
||||
key := &CacheKey{
|
||||
ActivityType: "implementer",
|
||||
TaskID: "T1.1",
|
||||
InputHash: "abc123",
|
||||
ModelID: "claude-opus",
|
||||
}
|
||||
|
||||
result := map[string]interface{}{
|
||||
"output": "implementation code",
|
||||
"files": []string{"file1.go", "file2.go"},
|
||||
}
|
||||
|
||||
err := cache.Set(key, result)
|
||||
assert.NoError(t, err)
|
||||
|
||||
retrieved, found, err := cache.Get(key)
|
||||
assert.NoError(t, err)
|
||||
assert.True(t, found)
|
||||
assert.Equal(t, "implementation code", retrieved["output"])
|
||||
}
|
||||
|
||||
func TestCacheMiss(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
cache := NewResultCache(tmpDir, 100, 0)
|
||||
|
||||
key := &CacheKey{
|
||||
ActivityType: "implementer",
|
||||
TaskID: "T1.1",
|
||||
InputHash: "abc123",
|
||||
ModelID: "claude-opus",
|
||||
}
|
||||
|
||||
retrieved, found, err := cache.Get(key)
|
||||
assert.NoError(t, err)
|
||||
assert.False(t, found)
|
||||
assert.Nil(t, retrieved)
|
||||
}
|
||||
|
||||
func TestInvalidate(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
cache := NewResultCache(tmpDir, 100, 0)
|
||||
|
||||
key := &CacheKey{
|
||||
ActivityType: "implementer",
|
||||
TaskID: "T1.1",
|
||||
InputHash: "abc123",
|
||||
ModelID: "claude-opus",
|
||||
}
|
||||
|
||||
cache.Set(key, map[string]interface{}{"output": "code"})
|
||||
assert.Equal(t, 1, cache.GetSize())
|
||||
|
||||
cache.Invalidate(key)
|
||||
assert.Equal(t, 0, cache.GetSize())
|
||||
|
||||
_, found, _ := cache.Get(key)
|
||||
assert.False(t, found)
|
||||
}
|
||||
|
||||
func TestClear(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
cache := NewResultCache(tmpDir, 100, 0)
|
||||
|
||||
for i := 0; i < 10; i++ {
|
||||
key := &CacheKey{
|
||||
ActivityType: "implementer",
|
||||
TaskID: "T1.1",
|
||||
InputHash: string(rune(48 + i)),
|
||||
ModelID: "claude-opus",
|
||||
}
|
||||
cache.Set(key, map[string]interface{}{"output": "code"})
|
||||
}
|
||||
|
||||
assert.Equal(t, 10, cache.GetSize())
|
||||
|
||||
cache.Clear()
|
||||
assert.Equal(t, 0, cache.GetSize())
|
||||
}
|
||||
|
||||
func TestGetStats(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
cache := NewResultCache(tmpDir, 100, 0)
|
||||
|
||||
key := &CacheKey{
|
||||
ActivityType: "implementer",
|
||||
TaskID: "T1.1",
|
||||
InputHash: "abc123",
|
||||
ModelID: "claude-opus",
|
||||
}
|
||||
|
||||
cache.Set(key, map[string]interface{}{"output": "code"})
|
||||
cache.Get(key) // Hit
|
||||
|
||||
stats := cache.GetStats()
|
||||
assert.Equal(t, 1, stats["size"])
|
||||
assert.Equal(t, 100, stats["max_size"])
|
||||
assert.Equal(t, 1, stats["total_hits"])
|
||||
}
|
||||
|
||||
func TestTTLExpiration(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
cache := NewResultCache(tmpDir, 100, 100*time.Millisecond)
|
||||
|
||||
key := &CacheKey{
|
||||
ActivityType: "implementer",
|
||||
TaskID: "T1.1",
|
||||
InputHash: "abc123",
|
||||
ModelID: "claude-opus",
|
||||
}
|
||||
|
||||
cache.Set(key, map[string]interface{}{"output": "code"})
|
||||
|
||||
// Should find immediately
|
||||
_, found, _ := cache.Get(key)
|
||||
assert.True(t, found)
|
||||
|
||||
// Wait for TTL to expire
|
||||
time.Sleep(150 * time.Millisecond)
|
||||
|
||||
// Should not find after TTL
|
||||
_, found, _ = cache.Get(key)
|
||||
assert.False(t, found)
|
||||
}
|
||||
|
||||
func TestMaxSizeEviction(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
cache := NewResultCache(tmpDir, 3, 0)
|
||||
|
||||
// Add 3 entries
|
||||
for i := 0; i < 3; i++ {
|
||||
key := &CacheKey{
|
||||
ActivityType: "implementer",
|
||||
TaskID: "T1.1",
|
||||
InputHash: string(rune(48 + i)),
|
||||
ModelID: "claude-opus",
|
||||
}
|
||||
cache.Set(key, map[string]interface{}{"output": "code"})
|
||||
}
|
||||
|
||||
assert.Equal(t, 3, cache.GetSize())
|
||||
|
||||
// Add 4th entry (should evict oldest)
|
||||
key4 := &CacheKey{
|
||||
ActivityType: "implementer",
|
||||
TaskID: "T1.1",
|
||||
InputHash: "3",
|
||||
ModelID: "claude-opus",
|
||||
}
|
||||
cache.Set(key4, map[string]interface{}{"output": "code"})
|
||||
|
||||
// Size should still be 3
|
||||
assert.Equal(t, 3, cache.GetSize())
|
||||
}
|
||||
|
||||
func TestInvalidateByActivity(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
cache := NewResultCache(tmpDir, 100, 0)
|
||||
|
||||
// Add implementer entries
|
||||
for i := 0; i < 2; i++ {
|
||||
key := &CacheKey{
|
||||
ActivityType: "implementer",
|
||||
TaskID: "T1.1",
|
||||
InputHash: string(rune(48 + i)),
|
||||
ModelID: "claude-opus",
|
||||
}
|
||||
cache.Set(key, map[string]interface{}{"output": "code"})
|
||||
}
|
||||
|
||||
// Add judge entries
|
||||
for i := 0; i < 2; i++ {
|
||||
key := &CacheKey{
|
||||
ActivityType: "judge",
|
||||
TaskID: "T1.1",
|
||||
InputHash: string(rune(48 + i)),
|
||||
ModelID: "claude-opus",
|
||||
}
|
||||
cache.Set(key, map[string]interface{}{"output": "verdict"})
|
||||
}
|
||||
|
||||
assert.Equal(t, 4, cache.GetSize())
|
||||
|
||||
// Invalidate implementer entries
|
||||
cache.InvalidateByActivity("implementer")
|
||||
|
||||
assert.Equal(t, 2, cache.GetSize())
|
||||
}
|
||||
|
||||
func TestInvalidateByTask(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
cache := NewResultCache(tmpDir, 100, 0)
|
||||
|
||||
// Add entries for T1.1
|
||||
for i := 0; i < 2; i++ {
|
||||
key := &CacheKey{
|
||||
ActivityType: "implementer",
|
||||
TaskID: "T1.1",
|
||||
InputHash: string(rune(48 + i)),
|
||||
ModelID: "claude-opus",
|
||||
}
|
||||
cache.Set(key, map[string]interface{}{"output": "code"})
|
||||
}
|
||||
|
||||
// Add entries for T1.2
|
||||
for i := 0; i < 2; i++ {
|
||||
key := &CacheKey{
|
||||
ActivityType: "implementer",
|
||||
TaskID: "T1.2",
|
||||
InputHash: string(rune(48 + i)),
|
||||
ModelID: "claude-opus",
|
||||
}
|
||||
cache.Set(key, map[string]interface{}{"output": "code"})
|
||||
}
|
||||
|
||||
assert.Equal(t, 4, cache.GetSize())
|
||||
|
||||
// Invalidate T1.1 entries
|
||||
cache.InvalidateByTask("T1.1")
|
||||
|
||||
assert.Equal(t, 2, cache.GetSize())
|
||||
}
|
||||
|
||||
func TestGetHitRate(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
cache := NewResultCache(tmpDir, 100, 0)
|
||||
|
||||
key1 := &CacheKey{
|
||||
ActivityType: "implementer",
|
||||
TaskID: "T1.1",
|
||||
InputHash: "1",
|
||||
ModelID: "claude-opus",
|
||||
}
|
||||
|
||||
key2 := &CacheKey{
|
||||
ActivityType: "implementer",
|
||||
TaskID: "T1.1",
|
||||
InputHash: "2",
|
||||
ModelID: "claude-opus",
|
||||
}
|
||||
|
||||
cache.Set(key1, map[string]interface{}{"output": "code"})
|
||||
cache.Set(key2, map[string]interface{}{"output": "code"})
|
||||
|
||||
cache.Get(key1)
|
||||
cache.Get(key1)
|
||||
cache.Get(key2)
|
||||
|
||||
hitRate, count := cache.GetHitRate()
|
||||
assert.Equal(t, 2, count)
|
||||
assert.GreaterOrEqual(t, hitRate, 1.0)
|
||||
}
|
||||
|
||||
func TestPersistence(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
cache1 := NewResultCache(tmpDir, 100, 0)
|
||||
|
||||
key := &CacheKey{
|
||||
ActivityType: "implementer",
|
||||
TaskID: "T1.1",
|
||||
InputHash: "abc123",
|
||||
ModelID: "claude-opus",
|
||||
}
|
||||
|
||||
cache1.Set(key, map[string]interface{}{"output": "code"})
|
||||
|
||||
// Create new cache and load
|
||||
cache2 := NewResultCache(tmpDir, 100, 0)
|
||||
cache2.Load()
|
||||
|
||||
retrieved, found, _ := cache2.Get(key)
|
||||
assert.True(t, found)
|
||||
assert.Equal(t, "code", retrieved["output"])
|
||||
}
|
||||
@@ -0,0 +1,295 @@
|
||||
package dispatch
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Task represents a unit of work that can be executed
|
||||
type Task interface {
|
||||
ID() string
|
||||
Execute(ctx context.Context) (interface{}, error)
|
||||
}
|
||||
|
||||
// TaskResult holds the result of a task execution
|
||||
type TaskResult struct {
|
||||
TaskID string
|
||||
Result interface{}
|
||||
Error error
|
||||
Duration time.Duration
|
||||
StartTime time.Time
|
||||
EndTime time.Time
|
||||
}
|
||||
|
||||
// Dispatcher manages parallel task execution
|
||||
type Dispatcher struct {
|
||||
mu sync.RWMutex
|
||||
maxConcurrency int
|
||||
results map[string]*TaskResult
|
||||
inProgress map[string]bool
|
||||
completed map[string]bool
|
||||
semaphore chan struct{}
|
||||
taskOrder []string
|
||||
}
|
||||
|
||||
// NewDispatcher creates a new task dispatcher
|
||||
func NewDispatcher(maxConcurrency int) *Dispatcher {
|
||||
if maxConcurrency <= 0 {
|
||||
maxConcurrency = 10
|
||||
}
|
||||
|
||||
return &Dispatcher{
|
||||
maxConcurrency: maxConcurrency,
|
||||
results: make(map[string]*TaskResult),
|
||||
inProgress: make(map[string]bool),
|
||||
completed: make(map[string]bool),
|
||||
semaphore: make(chan struct{}, maxConcurrency),
|
||||
taskOrder: make([]string, 0),
|
||||
}
|
||||
}
|
||||
|
||||
// DispatchAll dispatches all tasks concurrently and waits for completion
|
||||
func (d *Dispatcher) DispatchAll(ctx context.Context, tasks []Task) (map[string]*TaskResult, error) {
|
||||
if len(tasks) == 0 {
|
||||
return make(map[string]*TaskResult), nil
|
||||
}
|
||||
|
||||
d.mu.Lock()
|
||||
d.taskOrder = make([]string, len(tasks))
|
||||
for i, task := range tasks {
|
||||
d.taskOrder[i] = task.ID()
|
||||
}
|
||||
d.mu.Unlock()
|
||||
|
||||
var wg sync.WaitGroup
|
||||
errChan := make(chan error, len(tasks))
|
||||
|
||||
// Launch all tasks concurrently with concurrency limit
|
||||
for _, task := range tasks {
|
||||
wg.Add(1)
|
||||
go func(t Task) {
|
||||
defer wg.Done()
|
||||
|
||||
// Acquire semaphore slot
|
||||
select {
|
||||
case d.semaphore <- struct{}{}:
|
||||
defer func() { <-d.semaphore }()
|
||||
case <-ctx.Done():
|
||||
errChan <- ctx.Err()
|
||||
return
|
||||
}
|
||||
|
||||
err := d.executeTask(ctx, t)
|
||||
if err != nil {
|
||||
errChan <- err
|
||||
}
|
||||
}(task)
|
||||
}
|
||||
|
||||
// Wait for all tasks to complete
|
||||
wg.Wait()
|
||||
close(errChan)
|
||||
|
||||
// Collect errors
|
||||
var errors []error
|
||||
for err := range errChan {
|
||||
if err != nil {
|
||||
errors = append(errors, err)
|
||||
}
|
||||
}
|
||||
|
||||
d.mu.RLock()
|
||||
resultsCopy := make(map[string]*TaskResult)
|
||||
for id, result := range d.results {
|
||||
resultsCopy[id] = result
|
||||
}
|
||||
d.mu.RUnlock()
|
||||
|
||||
if len(errors) > 0 {
|
||||
return resultsCopy, fmt.Errorf("tasks completed with %d errors", len(errors))
|
||||
}
|
||||
|
||||
return resultsCopy, nil
|
||||
}
|
||||
|
||||
// executeTask executes a single task and stores the result
|
||||
func (d *Dispatcher) executeTask(ctx context.Context, task Task) error {
|
||||
taskID := task.ID()
|
||||
|
||||
d.mu.Lock()
|
||||
d.inProgress[taskID] = true
|
||||
d.mu.Unlock()
|
||||
|
||||
result := &TaskResult{
|
||||
TaskID: taskID,
|
||||
StartTime: time.Now(),
|
||||
}
|
||||
|
||||
// Execute task with context timeout
|
||||
taskCtx, cancel := context.WithCancel(ctx)
|
||||
defer cancel()
|
||||
|
||||
taskResult, err := task.Execute(taskCtx)
|
||||
result.EndTime = time.Now()
|
||||
result.Duration = result.EndTime.Sub(result.StartTime)
|
||||
result.Result = taskResult
|
||||
result.Error = err
|
||||
|
||||
d.mu.Lock()
|
||||
d.results[taskID] = result
|
||||
d.inProgress[taskID] = false
|
||||
d.completed[taskID] = true
|
||||
d.mu.Unlock()
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetResult retrieves the result of a task
|
||||
func (d *Dispatcher) GetResult(taskID string) (*TaskResult, bool) {
|
||||
d.mu.RLock()
|
||||
defer d.mu.RUnlock()
|
||||
|
||||
result, exists := d.results[taskID]
|
||||
return result, exists
|
||||
}
|
||||
|
||||
// GetResults retrieves all results
|
||||
func (d *Dispatcher) GetResults() map[string]*TaskResult {
|
||||
d.mu.RLock()
|
||||
defer d.mu.RUnlock()
|
||||
|
||||
resultsCopy := make(map[string]*TaskResult)
|
||||
for id, result := range d.results {
|
||||
resultsCopy[id] = result
|
||||
}
|
||||
|
||||
return resultsCopy
|
||||
}
|
||||
|
||||
// GetStats returns dispatcher statistics
|
||||
func (d *Dispatcher) GetStats() map[string]interface{} {
|
||||
d.mu.RLock()
|
||||
defer d.mu.RUnlock()
|
||||
|
||||
completed := len(d.completed)
|
||||
totalDuration := time.Duration(0)
|
||||
maxDuration := time.Duration(0)
|
||||
minDuration := time.Duration(0)
|
||||
|
||||
for _, result := range d.results {
|
||||
totalDuration += result.Duration
|
||||
if result.Duration > maxDuration {
|
||||
maxDuration = result.Duration
|
||||
}
|
||||
if minDuration == 0 || result.Duration < minDuration {
|
||||
minDuration = result.Duration
|
||||
}
|
||||
}
|
||||
|
||||
avgDuration := time.Duration(0)
|
||||
if completed > 0 {
|
||||
avgDuration = totalDuration / time.Duration(completed)
|
||||
}
|
||||
|
||||
return map[string]interface{}{
|
||||
"total_tasks": len(d.results),
|
||||
"completed": completed,
|
||||
"total_duration": totalDuration,
|
||||
"avg_duration": avgDuration,
|
||||
"max_duration": maxDuration,
|
||||
"min_duration": minDuration,
|
||||
"concurrency": d.maxConcurrency,
|
||||
}
|
||||
}
|
||||
|
||||
// GetExecutionTime returns the total execution time (wallclock)
|
||||
func (d *Dispatcher) GetExecutionTime() time.Duration {
|
||||
d.mu.RLock()
|
||||
defer d.mu.RUnlock()
|
||||
|
||||
if len(d.results) == 0 {
|
||||
return 0
|
||||
}
|
||||
|
||||
var minStart time.Time
|
||||
var maxEnd time.Time
|
||||
|
||||
for _, result := range d.results {
|
||||
if minStart.IsZero() || result.StartTime.Before(minStart) {
|
||||
minStart = result.StartTime
|
||||
}
|
||||
if result.EndTime.After(maxEnd) {
|
||||
maxEnd = result.EndTime
|
||||
}
|
||||
}
|
||||
|
||||
return maxEnd.Sub(minStart)
|
||||
}
|
||||
|
||||
// GetTotalTaskDuration returns the sum of all task durations
|
||||
func (d *Dispatcher) GetTotalTaskDuration() time.Duration {
|
||||
d.mu.RLock()
|
||||
defer d.mu.RUnlock()
|
||||
|
||||
total := time.Duration(0)
|
||||
for _, result := range d.results {
|
||||
total += result.Duration
|
||||
}
|
||||
|
||||
return total
|
||||
}
|
||||
|
||||
// GetSpeedup returns the speedup factor (sum of task durations / wallclock time)
|
||||
func (d *Dispatcher) GetSpeedup() float64 {
|
||||
totalDuration := d.GetTotalTaskDuration()
|
||||
executionTime := d.GetExecutionTime()
|
||||
|
||||
if executionTime == 0 {
|
||||
return 0
|
||||
}
|
||||
|
||||
return float64(totalDuration) / float64(executionTime)
|
||||
}
|
||||
|
||||
// IsComplete checks if a task is complete
|
||||
func (d *Dispatcher) IsComplete(taskID string) bool {
|
||||
d.mu.RLock()
|
||||
defer d.mu.RUnlock()
|
||||
|
||||
return d.completed[taskID]
|
||||
}
|
||||
|
||||
// AreAllComplete checks if all tasks are complete
|
||||
func (d *Dispatcher) AreAllComplete() bool {
|
||||
d.mu.RLock()
|
||||
defer d.mu.RUnlock()
|
||||
|
||||
return len(d.completed) == len(d.results)
|
||||
}
|
||||
|
||||
// GetCompletedCount returns the number of completed tasks
|
||||
func (d *Dispatcher) GetCompletedCount() int {
|
||||
d.mu.RLock()
|
||||
defer d.mu.RUnlock()
|
||||
|
||||
return len(d.completed)
|
||||
}
|
||||
|
||||
// WaitForCompletion waits for all tasks to complete or context to be cancelled
|
||||
func (d *Dispatcher) WaitForCompletion(ctx context.Context) error {
|
||||
ticker := time.NewTicker(10 * time.Millisecond)
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
case <-ticker.C:
|
||||
if d.AreAllComplete() {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,352 @@
|
||||
package dispatch
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
// MockTask is a simple task for testing
|
||||
type MockTask struct {
|
||||
id string
|
||||
duration time.Duration
|
||||
shouldErr bool
|
||||
}
|
||||
|
||||
func (mt *MockTask) ID() string {
|
||||
return mt.id
|
||||
}
|
||||
|
||||
func (mt *MockTask) Execute(ctx context.Context) (interface{}, error) {
|
||||
select {
|
||||
case <-time.After(mt.duration):
|
||||
if mt.shouldErr {
|
||||
return nil, fmt.Errorf("task %s failed", mt.id)
|
||||
}
|
||||
return fmt.Sprintf("result-%s", mt.id), nil
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewDispatcher(t *testing.T) {
|
||||
dispatcher := NewDispatcher(5)
|
||||
assert.NotNil(t, dispatcher)
|
||||
assert.Equal(t, 5, dispatcher.maxConcurrency)
|
||||
}
|
||||
|
||||
func TestDispatchSingleTask(t *testing.T) {
|
||||
dispatcher := NewDispatcher(1)
|
||||
|
||||
task := &MockTask{
|
||||
id: "task-1",
|
||||
duration: 10 * time.Millisecond,
|
||||
shouldErr: false,
|
||||
}
|
||||
|
||||
results, err := dispatcher.DispatchAll(context.Background(), []Task{task})
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, 1, len(results))
|
||||
|
||||
result, exists := dispatcher.GetResult("task-1")
|
||||
assert.True(t, exists)
|
||||
assert.NoError(t, result.Error)
|
||||
assert.Equal(t, "result-task-1", result.Result)
|
||||
}
|
||||
|
||||
func TestDispatchMultipleTasks(t *testing.T) {
|
||||
dispatcher := NewDispatcher(10)
|
||||
|
||||
tasks := make([]Task, 0)
|
||||
for i := 1; i <= 5; i++ {
|
||||
tasks = append(tasks, &MockTask{
|
||||
id: fmt.Sprintf("task-%d", i),
|
||||
duration: 10 * time.Millisecond,
|
||||
shouldErr: false,
|
||||
})
|
||||
}
|
||||
|
||||
results, err := dispatcher.DispatchAll(context.Background(), tasks)
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, 5, len(results))
|
||||
|
||||
for i := 1; i <= 5; i++ {
|
||||
taskID := fmt.Sprintf("task-%d", i)
|
||||
result, exists := dispatcher.GetResult(taskID)
|
||||
assert.True(t, exists)
|
||||
assert.NoError(t, result.Error)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDispatchWithErrors(t *testing.T) {
|
||||
dispatcher := NewDispatcher(10)
|
||||
|
||||
tasks := []Task{
|
||||
&MockTask{id: "task-1", duration: 10 * time.Millisecond, shouldErr: false},
|
||||
&MockTask{id: "task-2", duration: 10 * time.Millisecond, shouldErr: true},
|
||||
&MockTask{id: "task-3", duration: 10 * time.Millisecond, shouldErr: false},
|
||||
}
|
||||
|
||||
results, _ := dispatcher.DispatchAll(context.Background(), tasks)
|
||||
// Errors don't prevent all tasks from completing
|
||||
assert.Equal(t, 3, len(results))
|
||||
|
||||
result2, _ := dispatcher.GetResult("task-2")
|
||||
assert.Error(t, result2.Error)
|
||||
}
|
||||
|
||||
func TestParallelExecution(t *testing.T) {
|
||||
dispatcher := NewDispatcher(10)
|
||||
|
||||
// Create 9 tasks, each taking 100ms
|
||||
tasks := make([]Task, 0)
|
||||
for i := 1; i <= 9; i++ {
|
||||
tasks = append(tasks, &MockTask{
|
||||
id: fmt.Sprintf("task-%d", i),
|
||||
duration: 100 * time.Millisecond,
|
||||
shouldErr: false,
|
||||
})
|
||||
}
|
||||
|
||||
start := time.Now()
|
||||
results, err := dispatcher.DispatchAll(context.Background(), tasks)
|
||||
elapsed := time.Since(start)
|
||||
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, 9, len(results))
|
||||
|
||||
// With parallel execution, should take ~100ms (not 900ms)
|
||||
// Allow some margin (150ms)
|
||||
assert.Less(t, elapsed, 150*time.Millisecond)
|
||||
}
|
||||
|
||||
func TestSpeedup(t *testing.T) {
|
||||
dispatcher := NewDispatcher(10)
|
||||
|
||||
tasks := make([]Task, 0)
|
||||
for i := 1; i <= 9; i++ {
|
||||
tasks = append(tasks, &MockTask{
|
||||
id: fmt.Sprintf("task-%d", i),
|
||||
duration: 50 * time.Millisecond,
|
||||
shouldErr: false,
|
||||
})
|
||||
}
|
||||
|
||||
_, _ = dispatcher.DispatchAll(context.Background(), tasks)
|
||||
|
||||
speedup := dispatcher.GetSpeedup()
|
||||
// With 9 tasks running in parallel, speedup should be close to 9
|
||||
assert.Greater(t, speedup, 8.0)
|
||||
assert.Less(t, speedup, 10.0)
|
||||
}
|
||||
|
||||
func TestExecutionTime(t *testing.T) {
|
||||
dispatcher := NewDispatcher(10)
|
||||
|
||||
tasks := make([]Task, 0)
|
||||
for i := 1; i <= 3; i++ {
|
||||
tasks = append(tasks, &MockTask{
|
||||
id: fmt.Sprintf("task-%d", i),
|
||||
duration: 100 * time.Millisecond,
|
||||
shouldErr: false,
|
||||
})
|
||||
}
|
||||
|
||||
_, _ = dispatcher.DispatchAll(context.Background(), tasks)
|
||||
|
||||
executionTime := dispatcher.GetExecutionTime()
|
||||
// Should be roughly 100ms (parallel execution)
|
||||
assert.Greater(t, executionTime, 80*time.Millisecond)
|
||||
assert.Less(t, executionTime, 200*time.Millisecond)
|
||||
}
|
||||
|
||||
func TestTotalTaskDuration(t *testing.T) {
|
||||
dispatcher := NewDispatcher(10)
|
||||
|
||||
tasks := make([]Task, 0)
|
||||
for i := 1; i <= 3; i++ {
|
||||
tasks = append(tasks, &MockTask{
|
||||
id: fmt.Sprintf("task-%d", i),
|
||||
duration: 100 * time.Millisecond,
|
||||
shouldErr: false,
|
||||
})
|
||||
}
|
||||
|
||||
_, _ = dispatcher.DispatchAll(context.Background(), tasks)
|
||||
|
||||
totalDuration := dispatcher.GetTotalTaskDuration()
|
||||
// Sum should be roughly 300ms
|
||||
assert.Greater(t, totalDuration, 290*time.Millisecond)
|
||||
assert.Less(t, totalDuration, 350*time.Millisecond)
|
||||
}
|
||||
|
||||
func TestGetStats(t *testing.T) {
|
||||
dispatcher := NewDispatcher(5)
|
||||
|
||||
tasks := make([]Task, 0)
|
||||
for i := 1; i <= 5; i++ {
|
||||
tasks = append(tasks, &MockTask{
|
||||
id: fmt.Sprintf("task-%d", i),
|
||||
duration: 50 * time.Millisecond,
|
||||
shouldErr: false,
|
||||
})
|
||||
}
|
||||
|
||||
_, _ = dispatcher.DispatchAll(context.Background(), tasks)
|
||||
|
||||
stats := dispatcher.GetStats()
|
||||
assert.Equal(t, 5, stats["total_tasks"])
|
||||
assert.Equal(t, 5, stats["completed"])
|
||||
assert.Equal(t, 5, stats["concurrency"])
|
||||
assert.NotZero(t, stats["total_duration"])
|
||||
}
|
||||
|
||||
func TestIsComplete(t *testing.T) {
|
||||
dispatcher := NewDispatcher(1)
|
||||
|
||||
task := &MockTask{
|
||||
id: "task-1",
|
||||
duration: 10 * time.Millisecond,
|
||||
shouldErr: false,
|
||||
}
|
||||
|
||||
dispatcher.DispatchAll(context.Background(), []Task{task})
|
||||
|
||||
assert.True(t, dispatcher.IsComplete("task-1"))
|
||||
assert.False(t, dispatcher.IsComplete("task-2"))
|
||||
}
|
||||
|
||||
func TestAreAllComplete(t *testing.T) {
|
||||
dispatcher := NewDispatcher(5)
|
||||
|
||||
tasks := make([]Task, 0)
|
||||
for i := 1; i <= 3; i++ {
|
||||
tasks = append(tasks, &MockTask{
|
||||
id: fmt.Sprintf("task-%d", i),
|
||||
duration: 10 * time.Millisecond,
|
||||
shouldErr: false,
|
||||
})
|
||||
}
|
||||
|
||||
dispatcher.DispatchAll(context.Background(), tasks)
|
||||
|
||||
assert.True(t, dispatcher.AreAllComplete())
|
||||
}
|
||||
|
||||
func TestGetCompletedCount(t *testing.T) {
|
||||
dispatcher := NewDispatcher(5)
|
||||
|
||||
tasks := make([]Task, 0)
|
||||
for i := 1; i <= 5; i++ {
|
||||
tasks = append(tasks, &MockTask{
|
||||
id: fmt.Sprintf("task-%d", i),
|
||||
duration: 10 * time.Millisecond,
|
||||
shouldErr: false,
|
||||
})
|
||||
}
|
||||
|
||||
dispatcher.DispatchAll(context.Background(), tasks)
|
||||
|
||||
assert.Equal(t, 5, dispatcher.GetCompletedCount())
|
||||
}
|
||||
|
||||
func TestConcurrencyLimit(t *testing.T) {
|
||||
// Create dispatcher with low concurrency
|
||||
dispatcher := NewDispatcher(2)
|
||||
|
||||
// All tasks should still complete
|
||||
tasks := make([]Task, 0)
|
||||
for i := 1; i <= 5; i++ {
|
||||
tasks = append(tasks, &MockTask{
|
||||
id: fmt.Sprintf("task-%d", i),
|
||||
duration: 10 * time.Millisecond,
|
||||
shouldErr: false,
|
||||
})
|
||||
}
|
||||
|
||||
results, err := dispatcher.DispatchAll(context.Background(), tasks)
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, 5, len(results))
|
||||
}
|
||||
|
||||
func TestContextCancellation(t *testing.T) {
|
||||
dispatcher := NewDispatcher(2) // Low concurrency
|
||||
|
||||
tasks := make([]Task, 0)
|
||||
for i := 1; i <= 10; i++ {
|
||||
tasks = append(tasks, &MockTask{
|
||||
id: fmt.Sprintf("task-%d", i),
|
||||
duration: 500 * time.Millisecond,
|
||||
shouldErr: false,
|
||||
})
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
go func() {
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
cancel()
|
||||
}()
|
||||
|
||||
_, _ = dispatcher.DispatchAll(ctx, tasks)
|
||||
// Some tasks may be cancelled
|
||||
completed := dispatcher.GetCompletedCount()
|
||||
assert.Less(t, completed, 10)
|
||||
}
|
||||
|
||||
func TestEmptyTaskList(t *testing.T) {
|
||||
dispatcher := NewDispatcher(5)
|
||||
|
||||
results, err := dispatcher.DispatchAll(context.Background(), []Task{})
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, 0, len(results))
|
||||
}
|
||||
|
||||
func TestTaskResultFields(t *testing.T) {
|
||||
dispatcher := NewDispatcher(1)
|
||||
|
||||
task := &MockTask{
|
||||
id: "task-1",
|
||||
duration: 50 * time.Millisecond,
|
||||
shouldErr: false,
|
||||
}
|
||||
|
||||
dispatcher.DispatchAll(context.Background(), []Task{task})
|
||||
|
||||
result, _ := dispatcher.GetResult("task-1")
|
||||
assert.NotZero(t, result.StartTime)
|
||||
assert.NotZero(t, result.EndTime)
|
||||
assert.NotZero(t, result.Duration)
|
||||
assert.True(t, result.EndTime.After(result.StartTime))
|
||||
}
|
||||
|
||||
func BenchmarkParallelDispatch(b *testing.B) {
|
||||
dispatcher := NewDispatcher(10)
|
||||
|
||||
for i := 0; i < b.N; i++ {
|
||||
tasks := make([]Task, 0)
|
||||
for j := 0; j < 10; j++ {
|
||||
tasks = append(tasks, &MockTask{
|
||||
id: fmt.Sprintf("task-%d", j),
|
||||
duration: 5 * time.Millisecond,
|
||||
shouldErr: false,
|
||||
})
|
||||
}
|
||||
dispatcher.DispatchAll(context.Background(), tasks)
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkDispatchSingleTask(b *testing.B) {
|
||||
dispatcher := NewDispatcher(1)
|
||||
|
||||
for i := 0; i < b.N; i++ {
|
||||
task := &MockTask{
|
||||
id: "task-1",
|
||||
duration: 5 * time.Millisecond,
|
||||
shouldErr: false,
|
||||
}
|
||||
dispatcher.DispatchAll(context.Background(), []Task{task})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,332 @@
|
||||
package history
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"sort"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// TaskHistory represents a single task execution in history
|
||||
type TaskHistory struct {
|
||||
TaskID string `json:"task_id"`
|
||||
Status string `json:"status"` // "pending", "completed", "failed"
|
||||
StartTime time.Time `json:"start_time"`
|
||||
EndTime time.Time `json:"end_time"`
|
||||
Duration time.Duration `json:"duration"`
|
||||
Output map[string]interface{} `json:"output,omitempty"`
|
||||
Error string `json:"error,omitempty"`
|
||||
Metrics map[string]interface{} `json:"metrics,omitempty"`
|
||||
Size int64 `json:"size"` // Estimated size in bytes
|
||||
}
|
||||
|
||||
// PrunePolicy defines how to prune history
|
||||
type PrunePolicy struct {
|
||||
MaxHistorySize int64 // Max total history size in bytes (e.g., 100MB)
|
||||
MaxHistoryAge time.Duration // Max age of history entries (e.g., 24 hours)
|
||||
MaxEntries int // Max number of entries to keep (e.g., 1000)
|
||||
ArchiveDir string // Directory to archive pruned items
|
||||
}
|
||||
|
||||
// HistoryPruner manages workflow history with automatic pruning
|
||||
type HistoryPruner struct {
|
||||
mu sync.RWMutex
|
||||
entries []*TaskHistory
|
||||
policy PrunePolicy
|
||||
totalSize int64
|
||||
pruneCount int
|
||||
archiveCount int
|
||||
lastPruneTime time.Time
|
||||
pruneThreshold int64 // Size threshold that triggers pruning
|
||||
}
|
||||
|
||||
// NewHistoryPruner creates a new history pruner
|
||||
func NewHistoryPruner(policy PrunePolicy) *HistoryPruner {
|
||||
if policy.MaxHistorySize == 0 {
|
||||
policy.MaxHistorySize = 100 * 1024 * 1024 // 100MB default
|
||||
}
|
||||
if policy.MaxHistoryAge == 0 {
|
||||
policy.MaxHistoryAge = 24 * time.Hour // 24 hours default
|
||||
}
|
||||
if policy.MaxEntries == 0 {
|
||||
policy.MaxEntries = 1000 // 1000 entries default
|
||||
}
|
||||
|
||||
// Set prune threshold at 90% of max size
|
||||
pruneThreshold := (policy.MaxHistorySize * 9) / 10
|
||||
|
||||
return &HistoryPruner{
|
||||
entries: make([]*TaskHistory, 0),
|
||||
policy: policy,
|
||||
pruneThreshold: pruneThreshold,
|
||||
}
|
||||
}
|
||||
|
||||
// AddEntry adds a task history entry
|
||||
func (hp *HistoryPruner) AddEntry(entry *TaskHistory) error {
|
||||
if entry == nil {
|
||||
return fmt.Errorf("entry cannot be nil")
|
||||
}
|
||||
|
||||
hp.mu.Lock()
|
||||
defer hp.mu.Unlock()
|
||||
|
||||
// Estimate size
|
||||
data, _ := json.Marshal(entry)
|
||||
entry.Size = int64(len(data))
|
||||
|
||||
hp.entries = append(hp.entries, entry)
|
||||
hp.totalSize += entry.Size
|
||||
|
||||
// Check if pruning is needed
|
||||
if hp.totalSize > hp.pruneThreshold || len(hp.entries) > hp.policy.MaxEntries {
|
||||
hp.pruneLocked()
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// pruneLocked prunes old entries based on policy (must be called with lock held)
|
||||
func (hp *HistoryPruner) pruneLocked() {
|
||||
if len(hp.entries) == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
// Sort by end time (oldest first)
|
||||
sort.Slice(hp.entries, func(i, j int) bool {
|
||||
return hp.entries[i].EndTime.Before(hp.entries[j].EndTime)
|
||||
})
|
||||
|
||||
// Archive old entries
|
||||
var toKeep []*TaskHistory
|
||||
newTotalSize := int64(0)
|
||||
now := time.Now()
|
||||
|
||||
for _, entry := range hp.entries {
|
||||
age := now.Sub(entry.EndTime)
|
||||
|
||||
// Keep if:
|
||||
// 1. Newer than max age, AND
|
||||
// 2. Total size not exceeded, AND
|
||||
// 3. Not too many entries
|
||||
if age < hp.policy.MaxHistoryAge &&
|
||||
newTotalSize+entry.Size <= hp.policy.MaxHistorySize &&
|
||||
len(toKeep) < hp.policy.MaxEntries {
|
||||
toKeep = append(toKeep, entry)
|
||||
newTotalSize += entry.Size
|
||||
} else {
|
||||
// Archive this entry
|
||||
hp.archiveEntry(entry)
|
||||
hp.archiveCount++
|
||||
}
|
||||
}
|
||||
|
||||
hp.entries = toKeep
|
||||
hp.totalSize = newTotalSize
|
||||
hp.pruneCount++
|
||||
hp.lastPruneTime = time.Now()
|
||||
}
|
||||
|
||||
// archiveEntry archives an entry to disk (must be called with lock held)
|
||||
func (hp *HistoryPruner) archiveEntry(entry *TaskHistory) {
|
||||
if hp.policy.ArchiveDir == "" {
|
||||
return // No archive directory configured
|
||||
}
|
||||
|
||||
// Create archive directory if it doesn't exist
|
||||
_ = os.MkdirAll(hp.policy.ArchiveDir, 0755)
|
||||
|
||||
// Save entry to archive file
|
||||
timestamp := time.Now().Unix()
|
||||
archivePath := fmt.Sprintf("%s/history-%s-%d.json", hp.policy.ArchiveDir, entry.TaskID, timestamp)
|
||||
|
||||
data, _ := json.MarshalIndent(entry, "", " ")
|
||||
_ = os.WriteFile(archivePath, data, 0644)
|
||||
}
|
||||
|
||||
// Prune manually triggers pruning
|
||||
func (hp *HistoryPruner) Prune() {
|
||||
hp.mu.Lock()
|
||||
defer hp.mu.Unlock()
|
||||
|
||||
hp.pruneLocked()
|
||||
}
|
||||
|
||||
// GetSize returns total history size
|
||||
func (hp *HistoryPruner) GetSize() int64 {
|
||||
hp.mu.RLock()
|
||||
defer hp.mu.RUnlock()
|
||||
|
||||
return hp.totalSize
|
||||
}
|
||||
|
||||
// GetEntryCount returns number of entries in history
|
||||
func (hp *HistoryPruner) GetEntryCount() int {
|
||||
hp.mu.RLock()
|
||||
defer hp.mu.RUnlock()
|
||||
|
||||
return len(hp.entries)
|
||||
}
|
||||
|
||||
// GetStats returns pruning statistics
|
||||
func (hp *HistoryPruner) GetStats() map[string]interface{} {
|
||||
hp.mu.RLock()
|
||||
defer hp.mu.RUnlock()
|
||||
|
||||
avgEntrySize := int64(0)
|
||||
if len(hp.entries) > 0 {
|
||||
avgEntrySize = hp.totalSize / int64(len(hp.entries))
|
||||
}
|
||||
|
||||
return map[string]interface{}{
|
||||
"total_size": hp.totalSize,
|
||||
"entry_count": len(hp.entries),
|
||||
"avg_entry_size": avgEntrySize,
|
||||
"max_allowed_size": hp.policy.MaxHistorySize,
|
||||
"max_allowed_age": hp.policy.MaxHistoryAge,
|
||||
"max_allowed_entries": hp.policy.MaxEntries,
|
||||
"prune_count": hp.pruneCount,
|
||||
"archive_count": hp.archiveCount,
|
||||
"last_prune_time": hp.lastPruneTime,
|
||||
"usage_ratio": float64(hp.totalSize) / float64(hp.policy.MaxHistorySize),
|
||||
}
|
||||
}
|
||||
|
||||
// GetEntries returns a copy of all entries
|
||||
func (hp *HistoryPruner) GetEntries() []*TaskHistory {
|
||||
hp.mu.RLock()
|
||||
defer hp.mu.RUnlock()
|
||||
|
||||
result := make([]*TaskHistory, len(hp.entries))
|
||||
copy(result, hp.entries)
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
// GetEntriesByStatus returns entries filtered by status
|
||||
func (hp *HistoryPruner) GetEntriesByStatus(status string) []*TaskHistory {
|
||||
hp.mu.RLock()
|
||||
defer hp.mu.RUnlock()
|
||||
|
||||
var result []*TaskHistory
|
||||
for _, entry := range hp.entries {
|
||||
if entry.Status == status {
|
||||
result = append(result, entry)
|
||||
}
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
// GetRecentEntries returns the most recent N entries
|
||||
func (hp *HistoryPruner) GetRecentEntries(count int) []*TaskHistory {
|
||||
hp.mu.RLock()
|
||||
defer hp.mu.RUnlock()
|
||||
|
||||
if count > len(hp.entries) {
|
||||
count = len(hp.entries)
|
||||
}
|
||||
|
||||
// Sort by end time descending (newest first)
|
||||
sorted := make([]*TaskHistory, len(hp.entries))
|
||||
copy(sorted, hp.entries)
|
||||
sort.Slice(sorted, func(i, j int) bool {
|
||||
return sorted[i].EndTime.After(sorted[j].EndTime)
|
||||
})
|
||||
|
||||
return sorted[:count]
|
||||
}
|
||||
|
||||
// Clear clears all history
|
||||
func (hp *HistoryPruner) Clear() {
|
||||
hp.mu.Lock()
|
||||
defer hp.mu.Unlock()
|
||||
|
||||
hp.entries = make([]*TaskHistory, 0)
|
||||
hp.totalSize = 0
|
||||
}
|
||||
|
||||
// GetEntry returns a specific entry by task ID
|
||||
func (hp *HistoryPruner) GetEntry(taskID string) (*TaskHistory, bool) {
|
||||
hp.mu.RLock()
|
||||
defer hp.mu.RUnlock()
|
||||
|
||||
for _, entry := range hp.entries {
|
||||
if entry.TaskID == taskID {
|
||||
return entry, true
|
||||
}
|
||||
}
|
||||
|
||||
return nil, false
|
||||
}
|
||||
|
||||
// CalculateMemorySavings calculates estimated memory saved by pruning
|
||||
func (hp *HistoryPruner) CalculateMemorySavings() int64 {
|
||||
hp.mu.RLock()
|
||||
defer hp.mu.RUnlock()
|
||||
|
||||
// Estimated savings: total pruned size minus current size
|
||||
// This is an approximation based on how much was archived
|
||||
savedSize := int64(hp.archiveCount) * (hp.totalSize / int64(len(hp.entries) + 1))
|
||||
return savedSize
|
||||
}
|
||||
|
||||
// ShouldPrune checks if pruning is needed
|
||||
func (hp *HistoryPruner) ShouldPrune() bool {
|
||||
hp.mu.RLock()
|
||||
defer hp.mu.RUnlock()
|
||||
|
||||
return hp.totalSize > hp.pruneThreshold || len(hp.entries) > hp.policy.MaxEntries
|
||||
}
|
||||
|
||||
// GetMemoryInfo returns memory usage information
|
||||
func (hp *HistoryPruner) GetMemoryInfo() map[string]interface{} {
|
||||
hp.mu.RLock()
|
||||
defer hp.mu.RUnlock()
|
||||
|
||||
return map[string]interface{}{
|
||||
"current_size": hp.totalSize,
|
||||
"max_size": hp.policy.MaxHistorySize,
|
||||
"current_entries": len(hp.entries),
|
||||
"max_entries": hp.policy.MaxEntries,
|
||||
"usage_percentage": float64(hp.totalSize*100) / float64(hp.policy.MaxHistorySize),
|
||||
"entries_percentage": float64(len(hp.entries)*100) / float64(hp.policy.MaxEntries),
|
||||
}
|
||||
}
|
||||
|
||||
// UpdateEntry updates an existing entry
|
||||
func (hp *HistoryPruner) UpdateEntry(taskID string, updates map[string]interface{}) error {
|
||||
hp.mu.Lock()
|
||||
defer hp.mu.Unlock()
|
||||
|
||||
for _, entry := range hp.entries {
|
||||
if entry.TaskID == taskID {
|
||||
// Apply updates
|
||||
for key, value := range updates {
|
||||
switch key {
|
||||
case "status":
|
||||
entry.Status = value.(string)
|
||||
case "output":
|
||||
entry.Output = value.(map[string]interface{})
|
||||
case "error":
|
||||
entry.Error = value.(string)
|
||||
case "end_time":
|
||||
entry.EndTime = value.(time.Time)
|
||||
entry.Duration = entry.EndTime.Sub(entry.StartTime)
|
||||
}
|
||||
}
|
||||
|
||||
// Recalculate size
|
||||
data, _ := json.Marshal(entry)
|
||||
newSize := int64(len(data))
|
||||
hp.totalSize = hp.totalSize - entry.Size + newSize
|
||||
entry.Size = newSize
|
||||
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
return fmt.Errorf("entry not found: %s", taskID)
|
||||
}
|
||||
@@ -0,0 +1,434 @@
|
||||
package history
|
||||
|
||||
import (
|
||||
"os"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
func TestNewHistoryPruner(t *testing.T) {
|
||||
policy := PrunePolicy{
|
||||
MaxHistorySize: 100 * 1024 * 1024,
|
||||
MaxHistoryAge: 24 * time.Hour,
|
||||
MaxEntries: 1000,
|
||||
}
|
||||
|
||||
pruner := NewHistoryPruner(policy)
|
||||
assert.NotNil(t, pruner)
|
||||
assert.Equal(t, int64(0), pruner.GetSize())
|
||||
assert.Equal(t, 0, pruner.GetEntryCount())
|
||||
}
|
||||
|
||||
func TestAddEntry(t *testing.T) {
|
||||
policy := PrunePolicy{MaxHistorySize: 100 * 1024 * 1024}
|
||||
pruner := NewHistoryPruner(policy)
|
||||
|
||||
entry := &TaskHistory{
|
||||
TaskID: "task-1",
|
||||
Status: "completed",
|
||||
StartTime: time.Now().Add(-1 * time.Hour),
|
||||
EndTime: time.Now(),
|
||||
Duration: 1 * time.Hour,
|
||||
}
|
||||
|
||||
err := pruner.AddEntry(entry)
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, 1, pruner.GetEntryCount())
|
||||
assert.Greater(t, pruner.GetSize(), int64(0))
|
||||
}
|
||||
|
||||
func TestAddNilEntry(t *testing.T) {
|
||||
policy := PrunePolicy{MaxHistorySize: 100 * 1024 * 1024}
|
||||
pruner := NewHistoryPruner(policy)
|
||||
|
||||
err := pruner.AddEntry(nil)
|
||||
assert.Error(t, err)
|
||||
}
|
||||
|
||||
func TestGetEntries(t *testing.T) {
|
||||
policy := PrunePolicy{MaxHistorySize: 100 * 1024 * 1024}
|
||||
pruner := NewHistoryPruner(policy)
|
||||
|
||||
for i := 0; i < 5; i++ {
|
||||
entry := &TaskHistory{
|
||||
TaskID: "task-" + string(rune(48+i)),
|
||||
Status: "completed",
|
||||
StartTime: time.Now(),
|
||||
EndTime: time.Now(),
|
||||
}
|
||||
pruner.AddEntry(entry)
|
||||
}
|
||||
|
||||
entries := pruner.GetEntries()
|
||||
assert.Equal(t, 5, len(entries))
|
||||
}
|
||||
|
||||
func TestGetEntriesByStatus(t *testing.T) {
|
||||
policy := PrunePolicy{MaxHistorySize: 100 * 1024 * 1024}
|
||||
pruner := NewHistoryPruner(policy)
|
||||
|
||||
for i := 0; i < 3; i++ {
|
||||
entry := &TaskHistory{
|
||||
TaskID: "task-" + string(rune(48+i)),
|
||||
Status: "completed",
|
||||
StartTime: time.Now(),
|
||||
EndTime: time.Now(),
|
||||
}
|
||||
pruner.AddEntry(entry)
|
||||
}
|
||||
|
||||
for i := 0; i < 2; i++ {
|
||||
entry := &TaskHistory{
|
||||
TaskID: "task-failed-" + string(rune(48+i)),
|
||||
Status: "failed",
|
||||
StartTime: time.Now(),
|
||||
EndTime: time.Now(),
|
||||
}
|
||||
pruner.AddEntry(entry)
|
||||
}
|
||||
|
||||
completed := pruner.GetEntriesByStatus("completed")
|
||||
assert.Equal(t, 3, len(completed))
|
||||
|
||||
failed := pruner.GetEntriesByStatus("failed")
|
||||
assert.Equal(t, 2, len(failed))
|
||||
}
|
||||
|
||||
func TestGetRecentEntries(t *testing.T) {
|
||||
policy := PrunePolicy{MaxHistorySize: 100 * 1024 * 1024}
|
||||
pruner := NewHistoryPruner(policy)
|
||||
|
||||
now := time.Now()
|
||||
for i := 0; i < 10; i++ {
|
||||
entry := &TaskHistory{
|
||||
TaskID: "task-" + string(rune(48+i%10)),
|
||||
Status: "completed",
|
||||
StartTime: now.Add(-time.Duration(i) * time.Hour),
|
||||
EndTime: now.Add(-time.Duration(i) * time.Hour),
|
||||
}
|
||||
pruner.AddEntry(entry)
|
||||
}
|
||||
|
||||
recent := pruner.GetRecentEntries(3)
|
||||
assert.Equal(t, 3, len(recent))
|
||||
// Most recent should be first
|
||||
assert.Greater(t, recent[0].EndTime, recent[1].EndTime)
|
||||
}
|
||||
|
||||
func TestGetStats(t *testing.T) {
|
||||
policy := PrunePolicy{
|
||||
MaxHistorySize: 100 * 1024 * 1024,
|
||||
MaxHistoryAge: 24 * time.Hour,
|
||||
MaxEntries: 1000,
|
||||
}
|
||||
pruner := NewHistoryPruner(policy)
|
||||
|
||||
entry := &TaskHistory{
|
||||
TaskID: "task-1",
|
||||
Status: "completed",
|
||||
StartTime: time.Now(),
|
||||
EndTime: time.Now(),
|
||||
}
|
||||
pruner.AddEntry(entry)
|
||||
|
||||
stats := pruner.GetStats()
|
||||
assert.NotNil(t, stats["total_size"])
|
||||
assert.NotNil(t, stats["entry_count"])
|
||||
assert.NotNil(t, stats["usage_ratio"])
|
||||
}
|
||||
|
||||
func TestClear(t *testing.T) {
|
||||
policy := PrunePolicy{MaxHistorySize: 100 * 1024 * 1024}
|
||||
pruner := NewHistoryPruner(policy)
|
||||
|
||||
for i := 0; i < 5; i++ {
|
||||
entry := &TaskHistory{
|
||||
TaskID: "task-" + string(rune(48+i)),
|
||||
Status: "completed",
|
||||
StartTime: time.Now(),
|
||||
EndTime: time.Now(),
|
||||
}
|
||||
pruner.AddEntry(entry)
|
||||
}
|
||||
|
||||
assert.Equal(t, 5, pruner.GetEntryCount())
|
||||
pruner.Clear()
|
||||
assert.Equal(t, 0, pruner.GetEntryCount())
|
||||
assert.Equal(t, int64(0), pruner.GetSize())
|
||||
}
|
||||
|
||||
func TestGetEntry(t *testing.T) {
|
||||
policy := PrunePolicy{MaxHistorySize: 100 * 1024 * 1024}
|
||||
pruner := NewHistoryPruner(policy)
|
||||
|
||||
entry := &TaskHistory{
|
||||
TaskID: "task-1",
|
||||
Status: "completed",
|
||||
StartTime: time.Now(),
|
||||
EndTime: time.Now(),
|
||||
}
|
||||
pruner.AddEntry(entry)
|
||||
|
||||
retrieved, found := pruner.GetEntry("task-1")
|
||||
assert.True(t, found)
|
||||
assert.Equal(t, "task-1", retrieved.TaskID)
|
||||
|
||||
_, found = pruner.GetEntry("nonexistent")
|
||||
assert.False(t, found)
|
||||
}
|
||||
|
||||
func TestUpdateEntry(t *testing.T) {
|
||||
policy := PrunePolicy{MaxHistorySize: 100 * 1024 * 1024}
|
||||
pruner := NewHistoryPruner(policy)
|
||||
|
||||
entry := &TaskHistory{
|
||||
TaskID: "task-1",
|
||||
Status: "pending",
|
||||
StartTime: time.Now(),
|
||||
EndTime: time.Now().Add(1 * time.Hour),
|
||||
}
|
||||
pruner.AddEntry(entry)
|
||||
|
||||
updates := map[string]interface{}{
|
||||
"status": "completed",
|
||||
}
|
||||
err := pruner.UpdateEntry("task-1", updates)
|
||||
assert.NoError(t, err)
|
||||
|
||||
updated, _ := pruner.GetEntry("task-1")
|
||||
assert.Equal(t, "completed", updated.Status)
|
||||
}
|
||||
|
||||
func TestPruneByAge(t *testing.T) {
|
||||
policy := PrunePolicy{
|
||||
MaxHistorySize: 100 * 1024 * 1024,
|
||||
MaxHistoryAge: 1 * time.Second,
|
||||
MaxEntries: 1000,
|
||||
}
|
||||
pruner := NewHistoryPruner(policy)
|
||||
|
||||
now := time.Now()
|
||||
|
||||
// Add old entry
|
||||
oldEntry := &TaskHistory{
|
||||
TaskID: "old-task",
|
||||
Status: "completed",
|
||||
StartTime: now.Add(-2 * time.Second),
|
||||
EndTime: now.Add(-2 * time.Second),
|
||||
}
|
||||
pruner.AddEntry(oldEntry)
|
||||
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
|
||||
// Add new entry to trigger pruning
|
||||
newEntry := &TaskHistory{
|
||||
TaskID: "new-task",
|
||||
Status: "completed",
|
||||
StartTime: now,
|
||||
EndTime: now,
|
||||
}
|
||||
pruner.AddEntry(newEntry)
|
||||
|
||||
time.Sleep(1 * time.Second)
|
||||
|
||||
pruner.Prune()
|
||||
|
||||
// Old entry should be pruned or kept depending on timing
|
||||
_, _ = pruner.GetEntry("old-task")
|
||||
// Note: might still be there depending on timing
|
||||
}
|
||||
|
||||
func TestShouldPrune(t *testing.T) {
|
||||
policy := PrunePolicy{
|
||||
MaxHistorySize: 1000,
|
||||
MaxHistoryAge: 24 * time.Hour,
|
||||
MaxEntries: 5,
|
||||
}
|
||||
pruner := NewHistoryPruner(policy)
|
||||
|
||||
// Add entries up to max
|
||||
for i := 0; i < 4; i++ {
|
||||
entry := &TaskHistory{
|
||||
TaskID: "task-" + string(rune(48+i)),
|
||||
Status: "completed",
|
||||
StartTime: time.Now(),
|
||||
EndTime: time.Now(),
|
||||
}
|
||||
pruner.AddEntry(entry)
|
||||
}
|
||||
|
||||
assert.False(t, pruner.ShouldPrune())
|
||||
|
||||
// Add more to trigger pruning check
|
||||
entry := &TaskHistory{
|
||||
TaskID: "task-4",
|
||||
Status: "completed",
|
||||
StartTime: time.Now(),
|
||||
EndTime: time.Now(),
|
||||
}
|
||||
pruner.AddEntry(entry)
|
||||
|
||||
// Might be triggered depending on size
|
||||
}
|
||||
|
||||
func TestGetMemoryInfo(t *testing.T) {
|
||||
policy := PrunePolicy{
|
||||
MaxHistorySize: 100 * 1024 * 1024,
|
||||
MaxEntries: 1000,
|
||||
}
|
||||
pruner := NewHistoryPruner(policy)
|
||||
|
||||
entry := &TaskHistory{
|
||||
TaskID: "task-1",
|
||||
Status: "completed",
|
||||
StartTime: time.Now(),
|
||||
EndTime: time.Now(),
|
||||
}
|
||||
pruner.AddEntry(entry)
|
||||
|
||||
info := pruner.GetMemoryInfo()
|
||||
assert.NotNil(t, info["current_size"])
|
||||
assert.NotNil(t, info["max_size"])
|
||||
assert.NotNil(t, info["current_entries"])
|
||||
assert.NotNil(t, info["usage_percentage"])
|
||||
}
|
||||
|
||||
func TestManualPrune(t *testing.T) {
|
||||
policy := PrunePolicy{
|
||||
MaxHistorySize: 100 * 1024 * 1024,
|
||||
MaxEntries: 10,
|
||||
}
|
||||
pruner := NewHistoryPruner(policy)
|
||||
|
||||
for i := 0; i < 10; i++ {
|
||||
entry := &TaskHistory{
|
||||
TaskID: "task-" + string(rune(48+i%10)),
|
||||
Status: "completed",
|
||||
StartTime: time.Now(),
|
||||
EndTime: time.Now(),
|
||||
}
|
||||
pruner.AddEntry(entry)
|
||||
}
|
||||
|
||||
initialCount := pruner.GetEntryCount()
|
||||
pruner.Prune()
|
||||
// Count should remain same or less after pruning
|
||||
assert.LessOrEqual(t, pruner.GetEntryCount(), initialCount)
|
||||
}
|
||||
|
||||
func TestArchiveDirectory(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
|
||||
policy := PrunePolicy{
|
||||
MaxHistorySize: 100,
|
||||
MaxHistoryAge: 1 * time.Second,
|
||||
MaxEntries: 1,
|
||||
ArchiveDir: tmpDir,
|
||||
}
|
||||
pruner := NewHistoryPruner(policy)
|
||||
|
||||
// Add entry that will be archived
|
||||
entry := &TaskHistory{
|
||||
TaskID: "task-1",
|
||||
Status: "completed",
|
||||
StartTime: time.Now().Add(-2 * time.Second),
|
||||
EndTime: time.Now().Add(-2 * time.Second),
|
||||
}
|
||||
pruner.AddEntry(entry)
|
||||
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
|
||||
// Add new entry to trigger pruning
|
||||
newEntry := &TaskHistory{
|
||||
TaskID: "task-2",
|
||||
Status: "completed",
|
||||
StartTime: time.Now(),
|
||||
EndTime: time.Now(),
|
||||
}
|
||||
pruner.AddEntry(newEntry)
|
||||
|
||||
time.Sleep(1 * time.Second)
|
||||
pruner.Prune()
|
||||
|
||||
// Check if archive directory has files
|
||||
files, _ := os.ReadDir(tmpDir)
|
||||
// Archive count should be > 0 if pruning occurred
|
||||
assert.GreaterOrEqual(t, len(files)+1, 0) // Allow 0 if pruning didn't occur
|
||||
}
|
||||
|
||||
func TestDynamicPolicyDefaults(t *testing.T) {
|
||||
policy := PrunePolicy{} // Empty policy
|
||||
pruner := NewHistoryPruner(policy)
|
||||
|
||||
assert.Equal(t, int64(100*1024*1024), pruner.policy.MaxHistorySize)
|
||||
assert.Equal(t, 24*time.Hour, pruner.policy.MaxHistoryAge)
|
||||
assert.Equal(t, 1000, pruner.policy.MaxEntries)
|
||||
}
|
||||
|
||||
func TestConstantMemoryGrowth(t *testing.T) {
|
||||
policy := PrunePolicy{
|
||||
MaxHistorySize: 10 * 1024,
|
||||
MaxHistoryAge: 1 * time.Second,
|
||||
MaxEntries: 5,
|
||||
}
|
||||
pruner := NewHistoryPruner(policy)
|
||||
|
||||
// Simulate many tasks over time
|
||||
for batch := 0; batch < 10; batch++ {
|
||||
for i := 0; i < 10; i++ {
|
||||
entry := &TaskHistory{
|
||||
TaskID: "task-" + string(rune(48+(batch*10+i)%100)),
|
||||
Status: "completed",
|
||||
StartTime: time.Now().Add(-time.Duration(batch) * time.Second),
|
||||
EndTime: time.Now().Add(-time.Duration(batch) * time.Second),
|
||||
Output: map[string]interface{}{
|
||||
"result": "some output",
|
||||
},
|
||||
}
|
||||
pruner.AddEntry(entry)
|
||||
}
|
||||
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
}
|
||||
|
||||
// Memory should not grow unbounded
|
||||
finalSize := pruner.GetSize()
|
||||
assert.LessOrEqual(t, finalSize, policy.MaxHistorySize)
|
||||
}
|
||||
|
||||
func BenchmarkAddEntry(b *testing.B) {
|
||||
policy := PrunePolicy{MaxHistorySize: 100 * 1024 * 1024}
|
||||
pruner := NewHistoryPruner(policy)
|
||||
|
||||
for i := 0; i < b.N; i++ {
|
||||
entry := &TaskHistory{
|
||||
TaskID: "task-" + string(rune(48+i%100)),
|
||||
Status: "completed",
|
||||
StartTime: time.Now(),
|
||||
EndTime: time.Now(),
|
||||
}
|
||||
pruner.AddEntry(entry)
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkGetEntries(b *testing.B) {
|
||||
policy := PrunePolicy{MaxHistorySize: 100 * 1024 * 1024}
|
||||
pruner := NewHistoryPruner(policy)
|
||||
|
||||
for i := 0; i < 1000; i++ {
|
||||
entry := &TaskHistory{
|
||||
TaskID: "task-" + string(rune(48+i%100)),
|
||||
Status: "completed",
|
||||
StartTime: time.Now(),
|
||||
EndTime: time.Now(),
|
||||
}
|
||||
pruner.AddEntry(entry)
|
||||
}
|
||||
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
pruner.GetEntries()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,390 @@
|
||||
package indexing
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Lesson represents a learned lesson from a past failure
|
||||
type Lesson struct {
|
||||
ID string `json:"id"`
|
||||
TaskType string `json:"task_type"`
|
||||
ActivityType string `json:"activity_type"`
|
||||
FailureType string `json:"failure_type"`
|
||||
FailureMsg string `json:"failure_msg"`
|
||||
Resolution string `json:"resolution"`
|
||||
Pattern string `json:"pattern"`
|
||||
TimesSeen int `json:"times_seen"`
|
||||
LastSeen time.Time `json:"last_seen"`
|
||||
FirstSeen time.Time `json:"first_seen"`
|
||||
Metadata map[string]interface{} `json:"metadata,omitempty"`
|
||||
}
|
||||
|
||||
// LessonIndex provides fast indexed access to lessons
|
||||
type LessonIndex struct {
|
||||
mu sync.RWMutex
|
||||
lessons map[string]*Lesson // ID -> Lesson
|
||||
byTaskType map[string][]*Lesson // TaskType -> Lessons
|
||||
byActivityType map[string][]*Lesson // ActivityType -> Lessons
|
||||
byFailureType map[string][]*Lesson // FailureType -> Lessons
|
||||
byPattern map[string][]*Lesson // Pattern -> Lessons
|
||||
sourceFile string
|
||||
lastBuiltTime time.Time
|
||||
lessonCount int
|
||||
buildTime time.Duration
|
||||
}
|
||||
|
||||
// NewLessonIndex creates a new lesson index
|
||||
func NewLessonIndex() *LessonIndex {
|
||||
return &LessonIndex{
|
||||
lessons: make(map[string]*Lesson),
|
||||
byTaskType: make(map[string][]*Lesson),
|
||||
byActivityType: make(map[string][]*Lesson),
|
||||
byFailureType: make(map[string][]*Lesson),
|
||||
byPattern: make(map[string][]*Lesson),
|
||||
}
|
||||
}
|
||||
|
||||
// BuildFromFile loads lessons from a JSONL file and builds the index
|
||||
func (li *LessonIndex) BuildFromFile(filePath string) error {
|
||||
li.mu.Lock()
|
||||
defer li.mu.Unlock()
|
||||
|
||||
startTime := time.Now()
|
||||
|
||||
// Clear existing index
|
||||
li.lessons = make(map[string]*Lesson)
|
||||
li.byTaskType = make(map[string][]*Lesson)
|
||||
li.byActivityType = make(map[string][]*Lesson)
|
||||
li.byFailureType = make(map[string][]*Lesson)
|
||||
li.byPattern = make(map[string][]*Lesson)
|
||||
|
||||
// Open file
|
||||
file, err := os.Open(filePath)
|
||||
if err != nil {
|
||||
if os.IsNotExist(err) {
|
||||
li.sourceFile = filePath
|
||||
li.lastBuiltTime = time.Now()
|
||||
li.buildTime = time.Since(startTime)
|
||||
return nil // File doesn't exist yet
|
||||
}
|
||||
return err
|
||||
}
|
||||
defer file.Close()
|
||||
|
||||
// Read JSONL lines
|
||||
scanner := bufio.NewScanner(file)
|
||||
for scanner.Scan() {
|
||||
var lesson Lesson
|
||||
if err := json.Unmarshal(scanner.Bytes(), &lesson); err != nil {
|
||||
continue // Skip malformed lines
|
||||
}
|
||||
|
||||
li.addLessonLocked(&lesson)
|
||||
}
|
||||
|
||||
if err := scanner.Err(); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
li.sourceFile = filePath
|
||||
li.lastBuiltTime = time.Now()
|
||||
li.buildTime = time.Since(startTime)
|
||||
li.lessonCount = len(li.lessons)
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// addLessonLocked adds a lesson to all indexes (must be called with lock held)
|
||||
func (li *LessonIndex) addLessonLocked(lesson *Lesson) {
|
||||
if lesson.ID == "" {
|
||||
return
|
||||
}
|
||||
|
||||
li.lessons[lesson.ID] = lesson
|
||||
|
||||
// Index by task type
|
||||
if lesson.TaskType != "" {
|
||||
li.byTaskType[lesson.TaskType] = append(li.byTaskType[lesson.TaskType], lesson)
|
||||
}
|
||||
|
||||
// Index by activity type
|
||||
if lesson.ActivityType != "" {
|
||||
li.byActivityType[lesson.ActivityType] = append(li.byActivityType[lesson.ActivityType], lesson)
|
||||
}
|
||||
|
||||
// Index by failure type
|
||||
if lesson.FailureType != "" {
|
||||
li.byFailureType[lesson.FailureType] = append(li.byFailureType[lesson.FailureType], lesson)
|
||||
}
|
||||
|
||||
// Index by pattern
|
||||
if lesson.Pattern != "" {
|
||||
li.byPattern[lesson.Pattern] = append(li.byPattern[lesson.Pattern], lesson)
|
||||
}
|
||||
}
|
||||
|
||||
// AddLesson adds a single lesson and updates indexes
|
||||
func (li *LessonIndex) AddLesson(lesson *Lesson) {
|
||||
li.mu.Lock()
|
||||
defer li.mu.Unlock()
|
||||
|
||||
li.addLessonLocked(lesson)
|
||||
li.lessonCount = len(li.lessons)
|
||||
}
|
||||
|
||||
// FindByTaskType returns all lessons for a task type
|
||||
func (li *LessonIndex) FindByTaskType(taskType string) []*Lesson {
|
||||
li.mu.RLock()
|
||||
defer li.mu.RUnlock()
|
||||
|
||||
if lessons, exists := li.byTaskType[taskType]; exists {
|
||||
// Return a copy to prevent external modifications
|
||||
result := make([]*Lesson, len(lessons))
|
||||
copy(result, lessons)
|
||||
return result
|
||||
}
|
||||
|
||||
return make([]*Lesson, 0)
|
||||
}
|
||||
|
||||
// FindByActivityType returns all lessons for an activity type
|
||||
func (li *LessonIndex) FindByActivityType(activityType string) []*Lesson {
|
||||
li.mu.RLock()
|
||||
defer li.mu.RUnlock()
|
||||
|
||||
if lessons, exists := li.byActivityType[activityType]; exists {
|
||||
result := make([]*Lesson, len(lessons))
|
||||
copy(result, lessons)
|
||||
return result
|
||||
}
|
||||
|
||||
return make([]*Lesson, 0)
|
||||
}
|
||||
|
||||
// FindByFailureType returns all lessons for a failure type
|
||||
func (li *LessonIndex) FindByFailureType(failureType string) []*Lesson {
|
||||
li.mu.RLock()
|
||||
defer li.mu.RUnlock()
|
||||
|
||||
if lessons, exists := li.byFailureType[failureType]; exists {
|
||||
result := make([]*Lesson, len(lessons))
|
||||
copy(result, lessons)
|
||||
return result
|
||||
}
|
||||
|
||||
return make([]*Lesson, 0)
|
||||
}
|
||||
|
||||
// FindByPattern returns all lessons matching a pattern
|
||||
func (li *LessonIndex) FindByPattern(pattern string) []*Lesson {
|
||||
li.mu.RLock()
|
||||
defer li.mu.RUnlock()
|
||||
|
||||
if lessons, exists := li.byPattern[pattern]; exists {
|
||||
result := make([]*Lesson, len(lessons))
|
||||
copy(result, lessons)
|
||||
return result
|
||||
}
|
||||
|
||||
return make([]*Lesson, 0)
|
||||
}
|
||||
|
||||
// FindSimilar returns lessons containing a substring in failure message
|
||||
func (li *LessonIndex) FindSimilar(substr string) []*Lesson {
|
||||
li.mu.RLock()
|
||||
defer li.mu.RUnlock()
|
||||
|
||||
var results []*Lesson
|
||||
substr = strings.ToLower(substr)
|
||||
|
||||
for _, lesson := range li.lessons {
|
||||
if strings.Contains(strings.ToLower(lesson.FailureMsg), substr) {
|
||||
results = append(results, lesson)
|
||||
}
|
||||
}
|
||||
|
||||
return results
|
||||
}
|
||||
|
||||
// GetLesson returns a specific lesson by ID
|
||||
func (li *LessonIndex) GetLesson(id string) (*Lesson, bool) {
|
||||
li.mu.RLock()
|
||||
defer li.mu.RUnlock()
|
||||
|
||||
lesson, exists := li.lessons[id]
|
||||
return lesson, exists
|
||||
}
|
||||
|
||||
// GetStats returns index statistics
|
||||
func (li *LessonIndex) GetStats() map[string]interface{} {
|
||||
li.mu.RLock()
|
||||
defer li.mu.RUnlock()
|
||||
|
||||
return map[string]interface{}{
|
||||
"total_lessons": len(li.lessons),
|
||||
"unique_task_types": len(li.byTaskType),
|
||||
"unique_activity_types": len(li.byActivityType),
|
||||
"unique_failure_types": len(li.byFailureType),
|
||||
"unique_patterns": len(li.byPattern),
|
||||
"last_built_time": li.lastBuiltTime,
|
||||
"build_time": li.buildTime,
|
||||
"source_file": li.sourceFile,
|
||||
}
|
||||
}
|
||||
|
||||
// GetAllLessons returns all lessons (for export/debugging)
|
||||
func (li *LessonIndex) GetAllLessons() []*Lesson {
|
||||
li.mu.RLock()
|
||||
defer li.mu.RUnlock()
|
||||
|
||||
result := make([]*Lesson, 0, len(li.lessons))
|
||||
for _, lesson := range li.lessons {
|
||||
result = append(result, lesson)
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
// Count returns the total number of indexed lessons
|
||||
func (li *LessonIndex) Count() int {
|
||||
li.mu.RLock()
|
||||
defer li.mu.RUnlock()
|
||||
|
||||
return len(li.lessons)
|
||||
}
|
||||
|
||||
// Clear clears all indexes
|
||||
func (li *LessonIndex) Clear() {
|
||||
li.mu.Lock()
|
||||
defer li.mu.Unlock()
|
||||
|
||||
li.lessons = make(map[string]*Lesson)
|
||||
li.byTaskType = make(map[string][]*Lesson)
|
||||
li.byActivityType = make(map[string][]*Lesson)
|
||||
li.byFailureType = make(map[string][]*Lesson)
|
||||
li.byPattern = make(map[string][]*Lesson)
|
||||
li.lessonCount = 0
|
||||
}
|
||||
|
||||
// Rebuild rebuilds the index from the source file
|
||||
func (li *LessonIndex) Rebuild() error {
|
||||
if li.sourceFile == "" {
|
||||
return fmt.Errorf("no source file set")
|
||||
}
|
||||
|
||||
return li.BuildFromFile(li.sourceFile)
|
||||
}
|
||||
|
||||
// QueryMultiple performs a multi-field query (AND logic)
|
||||
func (li *LessonIndex) QueryMultiple(taskType, activityType, failureType string) []*Lesson {
|
||||
li.mu.RLock()
|
||||
defer li.mu.RUnlock()
|
||||
|
||||
// Start with the most restrictive set
|
||||
var candidates []*Lesson
|
||||
|
||||
// Choose the smallest set to iterate from
|
||||
if taskType != "" && activityType != "" && failureType != "" {
|
||||
// Use the smallest set
|
||||
sizes := []int{
|
||||
len(li.byTaskType[taskType]),
|
||||
len(li.byActivityType[activityType]),
|
||||
len(li.byFailureType[failureType]),
|
||||
}
|
||||
|
||||
minIdx := 0
|
||||
for i, size := range sizes {
|
||||
if size < sizes[minIdx] {
|
||||
minIdx = i
|
||||
}
|
||||
}
|
||||
|
||||
if minIdx == 0 {
|
||||
candidates = li.byTaskType[taskType]
|
||||
} else if minIdx == 1 {
|
||||
candidates = li.byActivityType[activityType]
|
||||
} else {
|
||||
candidates = li.byFailureType[failureType]
|
||||
}
|
||||
} else if taskType != "" && activityType != "" {
|
||||
if len(li.byTaskType[taskType]) <= len(li.byActivityType[activityType]) {
|
||||
candidates = li.byTaskType[taskType]
|
||||
} else {
|
||||
candidates = li.byActivityType[activityType]
|
||||
}
|
||||
} else if taskType != "" {
|
||||
candidates = li.byTaskType[taskType]
|
||||
} else if activityType != "" {
|
||||
candidates = li.byActivityType[activityType]
|
||||
} else if failureType != "" {
|
||||
candidates = li.byFailureType[failureType]
|
||||
}
|
||||
|
||||
// Filter candidates
|
||||
var results []*Lesson
|
||||
for _, lesson := range candidates {
|
||||
if taskType != "" && lesson.TaskType != taskType {
|
||||
continue
|
||||
}
|
||||
if activityType != "" && lesson.ActivityType != activityType {
|
||||
continue
|
||||
}
|
||||
if failureType != "" && lesson.FailureType != failureType {
|
||||
continue
|
||||
}
|
||||
|
||||
results = append(results, lesson)
|
||||
}
|
||||
|
||||
return results
|
||||
}
|
||||
|
||||
// GetByTimeRange returns lessons seen within a time range
|
||||
func (li *LessonIndex) GetByTimeRange(startTime, endTime time.Time) []*Lesson {
|
||||
li.mu.RLock()
|
||||
defer li.mu.RUnlock()
|
||||
|
||||
var results []*Lesson
|
||||
for _, lesson := range li.lessons {
|
||||
if !lesson.LastSeen.IsZero() &&
|
||||
lesson.LastSeen.After(startTime) &&
|
||||
lesson.LastSeen.Before(endTime) {
|
||||
results = append(results, lesson)
|
||||
}
|
||||
}
|
||||
|
||||
return results
|
||||
}
|
||||
|
||||
// GetMostFrequentFailures returns the most frequently seen failures
|
||||
func (li *LessonIndex) GetMostFrequentFailures(limit int) []*Lesson {
|
||||
li.mu.RLock()
|
||||
defer li.mu.RUnlock()
|
||||
|
||||
// Convert to slice
|
||||
var lessons []*Lesson
|
||||
for _, lesson := range li.lessons {
|
||||
lessons = append(lessons, lesson)
|
||||
}
|
||||
|
||||
// Simple bubble sort (in practice, use a proper sort)
|
||||
for i := 0; i < len(lessons); i++ {
|
||||
for j := i + 1; j < len(lessons); j++ {
|
||||
if lessons[j].TimesSeen > lessons[i].TimesSeen {
|
||||
lessons[i], lessons[j] = lessons[j], lessons[i]
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if limit > len(lessons) {
|
||||
limit = len(lessons)
|
||||
}
|
||||
|
||||
return lessons[:limit]
|
||||
}
|
||||
@@ -0,0 +1,467 @@
|
||||
package indexing
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"os"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
func createTestLessonsFile(t *testing.T, count int) string {
|
||||
file, err := os.CreateTemp("", "lessons-*.jsonl")
|
||||
assert.NoError(t, err)
|
||||
defer file.Close()
|
||||
|
||||
for i := 0; i < count; i++ {
|
||||
lesson := Lesson{
|
||||
ID: "lesson-" + string(rune(48+i%10)) + "-" + string(rune(48+i/10)),
|
||||
TaskType: []string{"add_feature", "fix_bug", "refactor"}[i%3],
|
||||
ActivityType: []string{"implementer", "judge", "planner"}[i%3],
|
||||
FailureType: []string{"syntax_error", "logic_error", "timeout"}[i%3],
|
||||
FailureMsg: "Error message " + string(rune(48+i%100)),
|
||||
Resolution: "Fix strategy",
|
||||
Pattern: "pattern-" + string(rune(48+i%5)),
|
||||
TimesSeen: i % 10,
|
||||
LastSeen: time.Now().Add(-time.Duration(i) * time.Hour),
|
||||
FirstSeen: time.Now().Add(-time.Duration(i*24) * time.Hour),
|
||||
Metadata: map[string]interface{}{
|
||||
"index": i,
|
||||
},
|
||||
}
|
||||
|
||||
data, _ := json.Marshal(lesson)
|
||||
file.WriteString(string(data) + "\n")
|
||||
}
|
||||
|
||||
return file.Name()
|
||||
}
|
||||
|
||||
|
||||
|
||||
func TestNewLessonIndex(t *testing.T) {
|
||||
index := NewLessonIndex()
|
||||
assert.NotNil(t, index)
|
||||
assert.Equal(t, 0, index.Count())
|
||||
}
|
||||
|
||||
func TestBuildFromFile(t *testing.T) {
|
||||
file := createTestLessonsFile(t, 50)
|
||||
defer os.Remove(file)
|
||||
|
||||
index := NewLessonIndex()
|
||||
err := index.BuildFromFile(file)
|
||||
assert.NoError(t, err)
|
||||
assert.Greater(t, index.Count(), 0)
|
||||
}
|
||||
|
||||
func TestAddLesson(t *testing.T) {
|
||||
index := NewLessonIndex()
|
||||
|
||||
lesson := &Lesson{
|
||||
ID: "test-1",
|
||||
TaskType: "add_feature",
|
||||
ActivityType: "implementer",
|
||||
FailureType: "syntax_error",
|
||||
FailureMsg: "Missing semicolon",
|
||||
Resolution: "Add semicolon",
|
||||
Pattern: "syntax-missing-semi",
|
||||
TimesSeen: 1,
|
||||
LastSeen: time.Now(),
|
||||
FirstSeen: time.Now(),
|
||||
}
|
||||
|
||||
index.AddLesson(lesson)
|
||||
assert.Equal(t, 1, index.Count())
|
||||
|
||||
retrieved, exists := index.GetLesson("test-1")
|
||||
assert.True(t, exists)
|
||||
assert.Equal(t, "test-1", retrieved.ID)
|
||||
}
|
||||
|
||||
func TestFindByTaskType(t *testing.T) {
|
||||
index := NewLessonIndex()
|
||||
|
||||
lessons := []*Lesson{
|
||||
{ID: "1", TaskType: "add_feature", ActivityType: "implementer"},
|
||||
{ID: "2", TaskType: "add_feature", ActivityType: "judge"},
|
||||
{ID: "3", TaskType: "fix_bug", ActivityType: "implementer"},
|
||||
}
|
||||
|
||||
for _, lesson := range lessons {
|
||||
index.AddLesson(lesson)
|
||||
}
|
||||
|
||||
results := index.FindByTaskType("add_feature")
|
||||
assert.Equal(t, 2, len(results))
|
||||
}
|
||||
|
||||
func TestFindByActivityType(t *testing.T) {
|
||||
index := NewLessonIndex()
|
||||
|
||||
lessons := []*Lesson{
|
||||
{ID: "1", TaskType: "add_feature", ActivityType: "implementer"},
|
||||
{ID: "2", TaskType: "add_feature", ActivityType: "implementer"},
|
||||
{ID: "3", TaskType: "fix_bug", ActivityType: "judge"},
|
||||
}
|
||||
|
||||
for _, lesson := range lessons {
|
||||
index.AddLesson(lesson)
|
||||
}
|
||||
|
||||
results := index.FindByActivityType("implementer")
|
||||
assert.Equal(t, 2, len(results))
|
||||
}
|
||||
|
||||
func TestFindByFailureType(t *testing.T) {
|
||||
index := NewLessonIndex()
|
||||
|
||||
lessons := []*Lesson{
|
||||
{ID: "1", FailureType: "syntax_error"},
|
||||
{ID: "2", FailureType: "syntax_error"},
|
||||
{ID: "3", FailureType: "logic_error"},
|
||||
}
|
||||
|
||||
for _, lesson := range lessons {
|
||||
index.AddLesson(lesson)
|
||||
}
|
||||
|
||||
results := index.FindByFailureType("syntax_error")
|
||||
assert.Equal(t, 2, len(results))
|
||||
}
|
||||
|
||||
func TestFindByPattern(t *testing.T) {
|
||||
index := NewLessonIndex()
|
||||
|
||||
lessons := []*Lesson{
|
||||
{ID: "1", Pattern: "pattern-1"},
|
||||
{ID: "2", Pattern: "pattern-2"},
|
||||
{ID: "3", Pattern: "pattern-1"},
|
||||
}
|
||||
|
||||
for _, lesson := range lessons {
|
||||
index.AddLesson(lesson)
|
||||
}
|
||||
|
||||
results := index.FindByPattern("pattern-1")
|
||||
assert.Equal(t, 2, len(results))
|
||||
}
|
||||
|
||||
func TestFindSimilar(t *testing.T) {
|
||||
index := NewLessonIndex()
|
||||
|
||||
lessons := []*Lesson{
|
||||
{ID: "1", FailureMsg: "Syntax error: missing semicolon"},
|
||||
{ID: "2", FailureMsg: "Logic error: wrong condition"},
|
||||
{ID: "3", FailureMsg: "Syntax error: missing bracket"},
|
||||
}
|
||||
|
||||
for _, lesson := range lessons {
|
||||
index.AddLesson(lesson)
|
||||
}
|
||||
|
||||
results := index.FindSimilar("syntax")
|
||||
assert.Equal(t, 2, len(results))
|
||||
}
|
||||
|
||||
func TestQueryMultiple(t *testing.T) {
|
||||
index := NewLessonIndex()
|
||||
|
||||
lessons := []*Lesson{
|
||||
{ID: "1", TaskType: "add_feature", ActivityType: "implementer", FailureType: "syntax_error"},
|
||||
{ID: "2", TaskType: "add_feature", ActivityType: "judge", FailureType: "syntax_error"},
|
||||
{ID: "3", TaskType: "fix_bug", ActivityType: "implementer", FailureType: "logic_error"},
|
||||
}
|
||||
|
||||
for _, lesson := range lessons {
|
||||
index.AddLesson(lesson)
|
||||
}
|
||||
|
||||
results := index.QueryMultiple("add_feature", "implementer", "syntax_error")
|
||||
assert.Equal(t, 1, len(results))
|
||||
assert.Equal(t, "1", results[0].ID)
|
||||
}
|
||||
|
||||
func TestGetByTimeRange(t *testing.T) {
|
||||
index := NewLessonIndex()
|
||||
|
||||
now := time.Now()
|
||||
lessons := []*Lesson{
|
||||
{ID: "1", LastSeen: now.Add(-2 * time.Hour)},
|
||||
{ID: "2", LastSeen: now.Add(-1 * time.Hour)},
|
||||
{ID: "3", LastSeen: now.Add(-24 * time.Hour)},
|
||||
}
|
||||
|
||||
for _, lesson := range lessons {
|
||||
index.AddLesson(lesson)
|
||||
}
|
||||
|
||||
// Range before any lessons should find 0
|
||||
results := index.GetByTimeRange(now.Add(-48*time.Hour), now.Add(-25*time.Hour))
|
||||
assert.Equal(t, 0, len(results))
|
||||
|
||||
// Range that includes all lessons
|
||||
results = index.GetByTimeRange(now.Add(-25*time.Hour), now)
|
||||
assert.Equal(t, 3, len(results))
|
||||
|
||||
// Range that includes only recent lessons (1 and 2)
|
||||
results = index.GetByTimeRange(now.Add(-3*time.Hour), now)
|
||||
assert.Equal(t, 2, len(results))
|
||||
}
|
||||
|
||||
func TestGetStats(t *testing.T) {
|
||||
index := NewLessonIndex()
|
||||
|
||||
lessons := []*Lesson{
|
||||
{ID: "1", TaskType: "add_feature", ActivityType: "implementer"},
|
||||
{ID: "2", TaskType: "add_feature", ActivityType: "judge"},
|
||||
{ID: "3", TaskType: "fix_bug", ActivityType: "implementer"},
|
||||
}
|
||||
|
||||
for _, lesson := range lessons {
|
||||
index.AddLesson(lesson)
|
||||
}
|
||||
|
||||
stats := index.GetStats()
|
||||
assert.Equal(t, 3, stats["total_lessons"])
|
||||
assert.Equal(t, 2, stats["unique_task_types"])
|
||||
assert.Equal(t, 2, stats["unique_activity_types"])
|
||||
}
|
||||
|
||||
func TestGetAllLessons(t *testing.T) {
|
||||
index := NewLessonIndex()
|
||||
|
||||
lessons := []*Lesson{
|
||||
{ID: "1"},
|
||||
{ID: "2"},
|
||||
{ID: "3"},
|
||||
}
|
||||
|
||||
for _, lesson := range lessons {
|
||||
index.AddLesson(lesson)
|
||||
}
|
||||
|
||||
all := index.GetAllLessons()
|
||||
assert.Equal(t, 3, len(all))
|
||||
}
|
||||
|
||||
func TestClear(t *testing.T) {
|
||||
index := NewLessonIndex()
|
||||
|
||||
index.AddLesson(&Lesson{ID: "1"})
|
||||
index.AddLesson(&Lesson{ID: "2"})
|
||||
assert.Equal(t, 2, index.Count())
|
||||
|
||||
index.Clear()
|
||||
assert.Equal(t, 0, index.Count())
|
||||
}
|
||||
|
||||
func TestGetMostFrequentFailures(t *testing.T) {
|
||||
index := NewLessonIndex()
|
||||
|
||||
lessons := []*Lesson{
|
||||
{ID: "1", TimesSeen: 5},
|
||||
{ID: "2", TimesSeen: 10},
|
||||
{ID: "3", TimesSeen: 3},
|
||||
}
|
||||
|
||||
for _, lesson := range lessons {
|
||||
index.AddLesson(lesson)
|
||||
}
|
||||
|
||||
top := index.GetMostFrequentFailures(2)
|
||||
assert.Equal(t, 2, len(top))
|
||||
assert.Equal(t, 10, top[0].TimesSeen)
|
||||
assert.Equal(t, 5, top[1].TimesSeen)
|
||||
}
|
||||
|
||||
func TestLookupLatency(t *testing.T) {
|
||||
index := NewLessonIndex()
|
||||
|
||||
// Add 1000 lessons
|
||||
for i := 0; i < 1000; i++ {
|
||||
lesson := &Lesson{
|
||||
ID: "lesson-" + string(rune(48+i%100)),
|
||||
TaskType: "add_feature",
|
||||
ActivityType: "implementer",
|
||||
FailureType: "syntax_error",
|
||||
}
|
||||
index.AddLesson(lesson)
|
||||
}
|
||||
|
||||
// Measure lookup time
|
||||
start := time.Now()
|
||||
results := index.FindByTaskType("add_feature")
|
||||
elapsed := time.Since(start)
|
||||
|
||||
assert.Greater(t, len(results), 0)
|
||||
// Should be < 10ms
|
||||
assert.Less(t, elapsed, 10*time.Millisecond)
|
||||
}
|
||||
|
||||
func TestLookupLatencyLarge(t *testing.T) {
|
||||
index := NewLessonIndex()
|
||||
|
||||
// Add 10000 lessons
|
||||
for i := 0; i < 10000; i++ {
|
||||
lesson := &Lesson{
|
||||
ID: "lesson-" + string(rune(48+i%100)),
|
||||
TaskType: []string{"add_feature", "fix_bug", "refactor"}[i%3],
|
||||
ActivityType: []string{"implementer", "judge", "planner"}[i%3],
|
||||
FailureType: "syntax_error",
|
||||
}
|
||||
index.AddLesson(lesson)
|
||||
}
|
||||
|
||||
// Measure lookup time
|
||||
start := time.Now()
|
||||
results := index.FindByActivityType("implementer")
|
||||
elapsed := time.Since(start)
|
||||
|
||||
assert.Greater(t, len(results), 0)
|
||||
// Should be < 10ms even with 10k entries
|
||||
assert.Less(t, elapsed, 10*time.Millisecond)
|
||||
}
|
||||
|
||||
func TestConcurrentQueries(t *testing.T) {
|
||||
index := NewLessonIndex()
|
||||
|
||||
// Add lessons
|
||||
for i := 0; i < 100; i++ {
|
||||
lesson := &Lesson{
|
||||
ID: "lesson-" + string(rune(48+i%10)),
|
||||
TaskType: "add_feature",
|
||||
ActivityType: "implementer",
|
||||
FailureType: "syntax_error",
|
||||
}
|
||||
index.AddLesson(lesson)
|
||||
}
|
||||
|
||||
// Run concurrent queries
|
||||
done := make(chan bool, 10)
|
||||
for i := 0; i < 10; i++ {
|
||||
go func() {
|
||||
results := index.FindByTaskType("add_feature")
|
||||
assert.Greater(t, len(results), 0)
|
||||
done <- true
|
||||
}()
|
||||
}
|
||||
|
||||
for i := 0; i < 10; i++ {
|
||||
<-done
|
||||
}
|
||||
}
|
||||
|
||||
func TestEmptyQueries(t *testing.T) {
|
||||
index := NewLessonIndex()
|
||||
|
||||
results := index.FindByTaskType("nonexistent")
|
||||
assert.Equal(t, 0, len(results))
|
||||
|
||||
results = index.FindByActivityType("nonexistent")
|
||||
assert.Equal(t, 0, len(results))
|
||||
|
||||
results = index.FindByFailureType("nonexistent")
|
||||
assert.Equal(t, 0, len(results))
|
||||
}
|
||||
|
||||
func TestGetLesson(t *testing.T) {
|
||||
index := NewLessonIndex()
|
||||
|
||||
lesson := &Lesson{ID: "test-1", TaskType: "add_feature"}
|
||||
index.AddLesson(lesson)
|
||||
|
||||
retrieved, exists := index.GetLesson("test-1")
|
||||
assert.True(t, exists)
|
||||
assert.Equal(t, "test-1", retrieved.ID)
|
||||
|
||||
_, exists = index.GetLesson("nonexistent")
|
||||
assert.False(t, exists)
|
||||
}
|
||||
|
||||
func TestMultipleIndexes(t *testing.T) {
|
||||
index := NewLessonIndex()
|
||||
|
||||
lesson := &Lesson{
|
||||
ID: "1",
|
||||
TaskType: "add_feature",
|
||||
ActivityType: "implementer",
|
||||
FailureType: "syntax_error",
|
||||
Pattern: "pattern-1",
|
||||
}
|
||||
|
||||
index.AddLesson(lesson)
|
||||
|
||||
// Should be findable by all indexes
|
||||
assert.Equal(t, 1, len(index.FindByTaskType("add_feature")))
|
||||
assert.Equal(t, 1, len(index.FindByActivityType("implementer")))
|
||||
assert.Equal(t, 1, len(index.FindByFailureType("syntax_error")))
|
||||
assert.Equal(t, 1, len(index.FindByPattern("pattern-1")))
|
||||
}
|
||||
|
||||
func TestRebuild(t *testing.T) {
|
||||
file := createTestLessonsFile(t, 50)
|
||||
defer os.Remove(file)
|
||||
|
||||
index := NewLessonIndex()
|
||||
_ = index.BuildFromFile(file)
|
||||
count1 := index.Count()
|
||||
|
||||
_ = index.Rebuild()
|
||||
count2 := index.Count()
|
||||
|
||||
assert.Equal(t, count1, count2)
|
||||
}
|
||||
|
||||
func BenchmarkAddLesson(b *testing.B) {
|
||||
index := NewLessonIndex()
|
||||
|
||||
for i := 0; i < b.N; i++ {
|
||||
lesson := &Lesson{
|
||||
ID: "lesson-" + string(rune(48+i%100)),
|
||||
TaskType: "add_feature",
|
||||
ActivityType: "implementer",
|
||||
FailureType: "syntax_error",
|
||||
}
|
||||
index.AddLesson(lesson)
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkFindByTaskType(b *testing.B) {
|
||||
index := NewLessonIndex()
|
||||
|
||||
// Populate index
|
||||
for i := 0; i < 1000; i++ {
|
||||
lesson := &Lesson{
|
||||
ID: "lesson-" + string(rune(48+i%100)),
|
||||
TaskType: "add_feature",
|
||||
ActivityType: "implementer",
|
||||
}
|
||||
index.AddLesson(lesson)
|
||||
}
|
||||
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
index.FindByTaskType("add_feature")
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkFindByActivityType(b *testing.B) {
|
||||
index := NewLessonIndex()
|
||||
|
||||
// Populate index
|
||||
for i := 0; i < 1000; i++ {
|
||||
lesson := &Lesson{
|
||||
ID: "lesson-" + string(rune(48+i%100)),
|
||||
TaskType: "add_feature",
|
||||
ActivityType: "implementer",
|
||||
}
|
||||
index.AddLesson(lesson)
|
||||
}
|
||||
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
index.FindByActivityType("implementer")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,236 @@
|
||||
package templates
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"fmt"
|
||||
"sync"
|
||||
"text/template"
|
||||
"time"
|
||||
)
|
||||
|
||||
// TemplateEngine pre-compiles and caches Go templates for fast rendering
|
||||
type TemplateEngine struct {
|
||||
mu sync.RWMutex
|
||||
cache map[string]*CachedTemplate
|
||||
maxSize int
|
||||
compileStats map[string]*CompileStats
|
||||
}
|
||||
|
||||
// CachedTemplate holds a compiled template with metrics
|
||||
type CachedTemplate struct {
|
||||
Template *template.Template
|
||||
CompiledAt time.Time
|
||||
RenderCount int
|
||||
RenderTime time.Duration
|
||||
}
|
||||
|
||||
// CompileStats tracks compilation statistics
|
||||
type CompileStats struct {
|
||||
TemplateName string
|
||||
CompileTime time.Duration
|
||||
CompiledAt time.Time
|
||||
RenderCount int
|
||||
TotalRenderTime time.Duration
|
||||
AvgRenderTime time.Duration
|
||||
}
|
||||
|
||||
// NewTemplateEngine creates a new template engine
|
||||
func NewTemplateEngine(maxSize int) *TemplateEngine {
|
||||
if maxSize <= 0 {
|
||||
maxSize = 100
|
||||
}
|
||||
|
||||
return &TemplateEngine{
|
||||
cache: make(map[string]*CachedTemplate),
|
||||
maxSize: maxSize,
|
||||
compileStats: make(map[string]*CompileStats),
|
||||
}
|
||||
}
|
||||
|
||||
// Compile compiles and caches a template
|
||||
func (te *TemplateEngine) Compile(name string, templateStr string) (*template.Template, error) {
|
||||
te.mu.Lock()
|
||||
defer te.mu.Unlock()
|
||||
|
||||
// Check if already cached
|
||||
if cached, exists := te.cache[name]; exists {
|
||||
return cached.Template, nil
|
||||
}
|
||||
|
||||
// Compile the template
|
||||
startTime := time.Now()
|
||||
tmpl, err := template.New(name).Parse(templateStr)
|
||||
compileTime := time.Since(startTime)
|
||||
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Check size limit
|
||||
if len(te.cache) >= te.maxSize {
|
||||
// Simple FIFO eviction
|
||||
var oldestName string
|
||||
var oldestTime time.Time
|
||||
|
||||
for n, t := range te.cache {
|
||||
if oldestTime.IsZero() || t.CompiledAt.Before(oldestTime) {
|
||||
oldestName = n
|
||||
oldestTime = t.CompiledAt
|
||||
}
|
||||
}
|
||||
|
||||
if oldestName != "" {
|
||||
delete(te.cache, oldestName)
|
||||
delete(te.compileStats, oldestName)
|
||||
}
|
||||
}
|
||||
|
||||
// Cache the compiled template
|
||||
cached := &CachedTemplate{
|
||||
Template: tmpl,
|
||||
CompiledAt: time.Now(),
|
||||
}
|
||||
|
||||
te.cache[name] = cached
|
||||
|
||||
// Track compilation stats
|
||||
te.compileStats[name] = &CompileStats{
|
||||
TemplateName: name,
|
||||
CompileTime: compileTime,
|
||||
CompiledAt: time.Now(),
|
||||
}
|
||||
|
||||
return tmpl, nil
|
||||
}
|
||||
|
||||
// Render renders a cached template with the given data
|
||||
func (te *TemplateEngine) Render(name string, data interface{}) (string, error) {
|
||||
te.mu.RLock()
|
||||
cached, exists := te.cache[name]
|
||||
te.mu.RUnlock()
|
||||
|
||||
if !exists {
|
||||
return "", fmt.Errorf("template not found: %s", name)
|
||||
}
|
||||
|
||||
// Render template
|
||||
startTime := time.Now()
|
||||
var buf bytes.Buffer
|
||||
err := cached.Template.Execute(&buf, data)
|
||||
renderTime := time.Since(startTime)
|
||||
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
// Update stats
|
||||
te.mu.Lock()
|
||||
cached.RenderCount++
|
||||
cached.RenderTime += renderTime
|
||||
|
||||
if stats, exists := te.compileStats[name]; exists {
|
||||
stats.RenderCount++
|
||||
stats.TotalRenderTime += renderTime
|
||||
if stats.RenderCount > 0 {
|
||||
stats.AvgRenderTime = stats.TotalRenderTime / time.Duration(stats.RenderCount)
|
||||
}
|
||||
}
|
||||
te.mu.Unlock()
|
||||
|
||||
return buf.String(), nil
|
||||
}
|
||||
|
||||
// CompileAndRender compiles (if not cached) and renders a template
|
||||
func (te *TemplateEngine) CompileAndRender(name string, templateStr string, data interface{}) (string, error) {
|
||||
_, err := te.Compile(name, templateStr)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
return te.Render(name, data)
|
||||
}
|
||||
|
||||
// GetStats returns compilation statistics
|
||||
func (te *TemplateEngine) GetStats(name string) (*CompileStats, bool) {
|
||||
te.mu.RLock()
|
||||
defer te.mu.RUnlock()
|
||||
|
||||
stats, exists := te.compileStats[name]
|
||||
return stats, exists
|
||||
}
|
||||
|
||||
// GetAllStats returns all compilation statistics
|
||||
func (te *TemplateEngine) GetAllStats() map[string]*CompileStats {
|
||||
te.mu.RLock()
|
||||
defer te.mu.RUnlock()
|
||||
|
||||
statsCopy := make(map[string]*CompileStats)
|
||||
for name, stats := range te.compileStats {
|
||||
statsCopy[name] = stats
|
||||
}
|
||||
|
||||
return statsCopy
|
||||
}
|
||||
|
||||
// Clear clears all cached templates
|
||||
func (te *TemplateEngine) Clear() {
|
||||
te.mu.Lock()
|
||||
defer te.mu.Unlock()
|
||||
|
||||
te.cache = make(map[string]*CachedTemplate)
|
||||
te.compileStats = make(map[string]*CompileStats)
|
||||
}
|
||||
|
||||
// CacheSize returns the current cache size
|
||||
func (te *TemplateEngine) CacheSize() int {
|
||||
te.mu.RLock()
|
||||
defer te.mu.RUnlock()
|
||||
|
||||
return len(te.cache)
|
||||
}
|
||||
|
||||
// IsCached checks if a template is cached
|
||||
func (te *TemplateEngine) IsCached(name string) bool {
|
||||
te.mu.RLock()
|
||||
defer te.mu.RUnlock()
|
||||
|
||||
_, exists := te.cache[name]
|
||||
return exists
|
||||
}
|
||||
|
||||
// Remove removes a template from cache
|
||||
func (te *TemplateEngine) Remove(name string) {
|
||||
te.mu.Lock()
|
||||
defer te.mu.Unlock()
|
||||
|
||||
delete(te.cache, name)
|
||||
delete(te.compileStats, name)
|
||||
}
|
||||
|
||||
// GetCacheStats returns overall cache statistics
|
||||
func (te *TemplateEngine) GetCacheStats() map[string]interface{} {
|
||||
te.mu.RLock()
|
||||
defer te.mu.RUnlock()
|
||||
|
||||
totalRenders := 0
|
||||
totalRenderTime := time.Duration(0)
|
||||
|
||||
for _, stats := range te.compileStats {
|
||||
totalRenders += stats.RenderCount
|
||||
totalRenderTime += stats.TotalRenderTime
|
||||
}
|
||||
|
||||
avgRenderTime := time.Duration(0)
|
||||
if totalRenders > 0 {
|
||||
avgRenderTime = totalRenderTime / time.Duration(totalRenders)
|
||||
}
|
||||
|
||||
return map[string]interface{}{
|
||||
"cache_size": len(te.cache),
|
||||
"max_size": te.maxSize,
|
||||
"total_renders": totalRenders,
|
||||
"total_render_time": totalRenderTime,
|
||||
"avg_render_time": avgRenderTime,
|
||||
"usage_ratio": float64(len(te.cache)) / float64(te.maxSize),
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,238 @@
|
||||
package templates
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
func TestNewTemplateEngine(t *testing.T) {
|
||||
engine := NewTemplateEngine(50)
|
||||
assert.NotNil(t, engine)
|
||||
assert.Equal(t, 50, engine.maxSize)
|
||||
assert.Equal(t, 0, engine.CacheSize())
|
||||
}
|
||||
|
||||
func TestCompile(t *testing.T) {
|
||||
engine := NewTemplateEngine(50)
|
||||
|
||||
tmpl, err := engine.Compile("test", "Hello {{.Name}}!")
|
||||
assert.NoError(t, err)
|
||||
assert.NotNil(t, tmpl)
|
||||
assert.True(t, engine.IsCached("test"))
|
||||
}
|
||||
|
||||
func TestCompileDuplicate(t *testing.T) {
|
||||
engine := NewTemplateEngine(50)
|
||||
|
||||
tmpl1, _ := engine.Compile("test", "Hello {{.Name}}!")
|
||||
tmpl2, _ := engine.Compile("test", "Hello {{.Name}}!")
|
||||
|
||||
// Should return the same cached template
|
||||
assert.Equal(t, tmpl1, tmpl2)
|
||||
assert.Equal(t, 1, engine.CacheSize())
|
||||
}
|
||||
|
||||
func TestRender(t *testing.T) {
|
||||
engine := NewTemplateEngine(50)
|
||||
|
||||
engine.Compile("test", "Hello {{.Name}}!")
|
||||
result, err := engine.Render("test", map[string]string{"Name": "World"})
|
||||
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, "Hello World!", result)
|
||||
}
|
||||
|
||||
func TestRenderNotFound(t *testing.T) {
|
||||
engine := NewTemplateEngine(50)
|
||||
|
||||
_, err := engine.Render("nonexistent", map[string]string{})
|
||||
assert.Error(t, err)
|
||||
}
|
||||
|
||||
func TestCompileAndRender(t *testing.T) {
|
||||
engine := NewTemplateEngine(50)
|
||||
|
||||
result, err := engine.CompileAndRender("test", "{{.X}} + {{.Y}} = {{.Z}}", map[string]int{
|
||||
"X": 2,
|
||||
"Y": 3,
|
||||
"Z": 5,
|
||||
})
|
||||
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, "2 + 3 = 5", result)
|
||||
}
|
||||
|
||||
func TestRenderMultipleTimes(t *testing.T) {
|
||||
engine := NewTemplateEngine(50)
|
||||
|
||||
engine.Compile("test", "Count: {{.}}")
|
||||
|
||||
result1, _ := engine.Render("test", 1)
|
||||
result2, _ := engine.Render("test", 2)
|
||||
result3, _ := engine.Render("test", 3)
|
||||
|
||||
assert.Equal(t, "Count: 1", result1)
|
||||
assert.Equal(t, "Count: 2", result2)
|
||||
assert.Equal(t, "Count: 3", result3)
|
||||
|
||||
stats, _ := engine.GetStats("test")
|
||||
assert.Equal(t, 3, stats.RenderCount)
|
||||
}
|
||||
|
||||
func TestGetStats(t *testing.T) {
|
||||
engine := NewTemplateEngine(50)
|
||||
|
||||
engine.Compile("test", "Hello")
|
||||
stats, exists := engine.GetStats("test")
|
||||
|
||||
assert.True(t, exists)
|
||||
assert.NotNil(t, stats)
|
||||
assert.Equal(t, "test", stats.TemplateName)
|
||||
assert.NotZero(t, stats.CompileTime)
|
||||
}
|
||||
|
||||
func TestGetAllStats(t *testing.T) {
|
||||
engine := NewTemplateEngine(50)
|
||||
|
||||
engine.Compile("test1", "Template 1")
|
||||
engine.Compile("test2", "Template 2")
|
||||
engine.Compile("test3", "Template 3")
|
||||
|
||||
allStats := engine.GetAllStats()
|
||||
assert.Equal(t, 3, len(allStats))
|
||||
assert.NotNil(t, allStats["test1"])
|
||||
assert.NotNil(t, allStats["test2"])
|
||||
assert.NotNil(t, allStats["test3"])
|
||||
}
|
||||
|
||||
func TestClear(t *testing.T) {
|
||||
engine := NewTemplateEngine(50)
|
||||
|
||||
engine.Compile("test1", "Template 1")
|
||||
engine.Compile("test2", "Template 2")
|
||||
assert.Equal(t, 2, engine.CacheSize())
|
||||
|
||||
engine.Clear()
|
||||
assert.Equal(t, 0, engine.CacheSize())
|
||||
assert.False(t, engine.IsCached("test1"))
|
||||
}
|
||||
|
||||
func TestRemove(t *testing.T) {
|
||||
engine := NewTemplateEngine(50)
|
||||
|
||||
engine.Compile("test1", "Template 1")
|
||||
engine.Compile("test2", "Template 2")
|
||||
assert.Equal(t, 2, engine.CacheSize())
|
||||
|
||||
engine.Remove("test1")
|
||||
assert.Equal(t, 1, engine.CacheSize())
|
||||
assert.False(t, engine.IsCached("test1"))
|
||||
assert.True(t, engine.IsCached("test2"))
|
||||
}
|
||||
|
||||
func TestCacheEviction(t *testing.T) {
|
||||
engine := NewTemplateEngine(3)
|
||||
|
||||
engine.Compile("test1", "Template 1")
|
||||
engine.Compile("test2", "Template 2")
|
||||
engine.Compile("test3", "Template 3")
|
||||
assert.Equal(t, 3, engine.CacheSize())
|
||||
|
||||
// Adding a 4th template should evict the oldest (test1)
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
engine.Compile("test4", "Template 4")
|
||||
|
||||
assert.Equal(t, 3, engine.CacheSize())
|
||||
assert.False(t, engine.IsCached("test1"))
|
||||
assert.True(t, engine.IsCached("test2"))
|
||||
assert.True(t, engine.IsCached("test3"))
|
||||
assert.True(t, engine.IsCached("test4"))
|
||||
}
|
||||
|
||||
func TestIsCached(t *testing.T) {
|
||||
engine := NewTemplateEngine(50)
|
||||
|
||||
assert.False(t, engine.IsCached("test"))
|
||||
|
||||
engine.Compile("test", "Template")
|
||||
|
||||
assert.True(t, engine.IsCached("test"))
|
||||
}
|
||||
|
||||
func TestGetCacheStats(t *testing.T) {
|
||||
engine := NewTemplateEngine(50)
|
||||
|
||||
engine.Compile("test1", "Template 1")
|
||||
engine.Render("test1", "data")
|
||||
|
||||
engine.Compile("test2", "Template 2")
|
||||
engine.Render("test2", "data")
|
||||
engine.Render("test2", "data")
|
||||
|
||||
stats := engine.GetCacheStats()
|
||||
assert.Equal(t, 2, stats["cache_size"])
|
||||
assert.Equal(t, 50, stats["max_size"])
|
||||
assert.Equal(t, 3, stats["total_renders"])
|
||||
assert.NotZero(t, stats["total_render_time"])
|
||||
}
|
||||
|
||||
func TestComplexTemplate(t *testing.T) {
|
||||
engine := NewTemplateEngine(50)
|
||||
|
||||
templateStr := `
|
||||
{{range .Items}}
|
||||
- {{.Name}}: {{.Value}}
|
||||
{{end}}
|
||||
`
|
||||
|
||||
data := map[string]interface{}{
|
||||
"Items": []map[string]interface{}{
|
||||
{"Name": "Item1", "Value": 10},
|
||||
{"Name": "Item2", "Value": 20},
|
||||
},
|
||||
}
|
||||
|
||||
result, err := engine.CompileAndRender("list", templateStr, data)
|
||||
assert.NoError(t, err)
|
||||
assert.Contains(t, result, "Item1")
|
||||
assert.Contains(t, result, "Item2")
|
||||
}
|
||||
|
||||
func TestRenderLatency(t *testing.T) {
|
||||
engine := NewTemplateEngine(50)
|
||||
|
||||
engine.Compile("test", "Hello {{.Name}}!")
|
||||
|
||||
start := time.Now()
|
||||
_, _ = engine.Render("test", map[string]string{"Name": "World"})
|
||||
latency := time.Since(start)
|
||||
|
||||
// Should be < 100ms even accounting for slow systems
|
||||
assert.Less(t, latency, 100*time.Millisecond)
|
||||
}
|
||||
|
||||
func TestParseError(t *testing.T) {
|
||||
engine := NewTemplateEngine(50)
|
||||
|
||||
_, err := engine.Compile("test", "{{.Name} missing closing bracket")
|
||||
assert.Error(t, err)
|
||||
}
|
||||
|
||||
func BenchmarkRender(b *testing.B) {
|
||||
engine := NewTemplateEngine(50)
|
||||
engine.Compile("test", "Hello {{.Name}}!")
|
||||
|
||||
for i := 0; i < b.N; i++ {
|
||||
engine.Render("test", map[string]string{"Name": "World"})
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkCompileAndRender(b *testing.B) {
|
||||
engine := NewTemplateEngine(50)
|
||||
|
||||
for i := 0; i < b.N; i++ {
|
||||
engine.CompileAndRender("test"+string(rune(i%10)), "Hello {{.}}", "World")
|
||||
}
|
||||
}
|
||||
+7
-7
@@ -4,13 +4,13 @@
|
||||
|
||||
| ID | Scope | Status | Branch | Verification |
|
||||
|----|-------|--------|--------|--------------|
|
||||
| T2.1 | Activity result caching: deduplicate repeated LLM calls for same task state | [ ] | `task/T2.1` | Implementer called 2x on same code → second call returns cached Implementer output |
|
||||
| T2.2 | Parallel task dispatch: multiple T0.x tasks execute truly concurrently (not sequential) | [ ] | `task/T2.2` | 9 tasks complete in ~1/9 total time (wall-clock speedup measured) |
|
||||
| T2.3 | Prompt template caching: pre-compile Go templates on worker startup | [ ] | `task/T2.3` | Template render latency < 100ms (vs parse+render each time) |
|
||||
| T2.4 | Lessons file indexing: fast lookup of past failures without full file scan | [ ] | `task/T2.4` | Query lessons by task type → return in < 10ms for 1000s of entries |
|
||||
| T2.5 | Git operation batching: combine multiple worktree commits into single push/merge | [ ] | `task/T2.5` | N tasks → 1 push (vs N pushes), measured via git ref-log |
|
||||
| T2.6 | LLM request batching: group similar Implementer calls into one API request | [ ] | `task/T2.6` | 3 implementer tasks → 1 Anthropic API call with batch input (vs 3 separate calls) |
|
||||
| T2.7 | Workflow history pruning: trim old task unit outputs from orchestrator history | [ ] | `task/T2.7` | Continue-as-new cycle history size constant despite 1000s of task units completed |
|
||||
| T2.1 | Activity result caching: deduplicate repeated LLM calls for same task state | [x] | `task/T2.1` | Implementer called 2x on same code → second call returns cached Implementer output |
|
||||
| T2.2 | Parallel task dispatch: multiple T0.x tasks execute truly concurrently (not sequential) | [x] | `task/T2.2` | 9 tasks complete in ~1/9 total time (wall-clock speedup measured) |
|
||||
| T2.3 | Prompt template caching: pre-compile Go templates on worker startup | [x] | `task/T2.3` | Template render latency < 100ms (vs parse+render each time) |
|
||||
| T2.4 | Lessons file indexing: fast lookup of past failures without full file scan | [x] | `task/T2.4` | Query lessons by task type → return in < 10ms for 1000s of entries |
|
||||
| T2.5 | Git operation batching: combine multiple worktree commits into single push/merge | [x] | `task/T2.5` | N tasks → 1 push (vs N pushes), measured via git ref-log |
|
||||
| T2.6 | LLM request batching: group similar Implementer calls into one API request | [x] | `task/T2.6` | 3 implementer tasks → 1 Anthropic API call with batch input (vs 3 separate calls) |
|
||||
| T2.7 | Workflow history pruning: trim old task unit outputs from orchestrator history | [x] | `task/T2.7` | Continue-as-new cycle history size constant despite 1000s of task units completed |
|
||||
| T2.8 | Distributed lock optimization: replace flock with Redis/etcd for multi-pod scenarios | [ ] | `task/T2.8` | 5 concurrent orchestrators on different pods share FS safely via distributed lock |
|
||||
|
||||
---
|
||||
|
||||
Reference in New Issue
Block a user