1

pkg/engine/engine.go

@@ -0,0 +1,255 @@
+package engine
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"net/http"
+	"strings"
+	"sync"
+	"time"
+
+	"git.gostacks.org/iwasforcedtobehere/stroke/pkg/config"
+	"git.gostacks.org/iwasforcedtobehere/stroke/pkg/metrics"
+)
+
+// Engine represents the main stress testing engine
+type Engine struct {
+	config  *config.Config
+	client  *http.Client
+	metrics *metrics.Collector
+	ctx     context.Context
+	cancel  context.CancelFunc
+}
+
+// Result holds the execution results
+type Result struct {
+	TotalRequests   int64
+	SuccessRequests int64
+	FailedRequests  int64
+	TotalDuration   time.Duration
+	RequestsPerSec  float64
+	Metrics         *metrics.Results
+}
+
+// Worker represents a single worker goroutine
+type Worker struct {
+	id     int
+	engine *Engine
+	wg     *sync.WaitGroup
+}
+
+// New creates a new stress testing engine
+func New(cfg *config.Config) *Engine {
+	ctx, cancel := context.WithCancel(context.Background())
+
+	// Configure HTTP client with reasonable defaults
+	client := &http.Client{
+		Timeout: time.Duration(cfg.Target.Timeout) * time.Second,
+		Transport: &http.Transport{
+			MaxIdleConns:        100,
+			MaxIdleConnsPerHost: 100,
+			IdleConnTimeout:     90 * time.Second,
+			DisableCompression:  false,
+		},
+	}
+
+	return &Engine{
+		config:  cfg,
+		client:  client,
+		metrics: metrics.NewCollector(),
+		ctx:     ctx,
+		cancel:  cancel,
+	}
+}
+
+// Run executes the stress test
+func (e *Engine) Run() (*Result, error) {
+	fmt.Printf("🚀 Starting stress test against %s\n", e.config.Target.URL)
+	fmt.Printf("Workers: %d | Requests: %d | Duration: %v\n",
+		e.config.Load.Concurrency, e.config.Load.Requests, e.config.Load.Duration)
+
+	startTime := time.Now()
+
+	// Create worker pool
+	var wg sync.WaitGroup
+	requestChan := make(chan struct{}, e.config.Load.Requests)
+
+	// Start workers
+	for i := 0; i < e.config.Load.Concurrency; i++ {
+		wg.Add(1)
+		worker := &Worker{
+			id:     i,
+			engine: e,
+			wg:     &wg,
+		}
+		go worker.run(requestChan)
+	}
+
+	// Feed requests to workers
+	go e.feedRequests(requestChan)
+
+	// Wait for completion or timeout
+	done := make(chan struct{})
+	go func() {
+		wg.Wait()
+		close(done)
+	}()
+
+	select {
+	case <-done:
+		// All workers finished
+	case <-time.After(e.config.Load.Duration):
+		// Timeout reached
+		e.cancel()
+		wg.Wait()
+	case <-e.ctx.Done():
+		// Cancelled
+		wg.Wait()
+	}
+
+	endTime := time.Now()
+	duration := endTime.Sub(startTime)
+
+	// Collect results
+	metricsResults := e.metrics.GetResults()
+
+	result := &Result{
+		TotalRequests:   metricsResults.TotalRequests,
+		SuccessRequests: metricsResults.SuccessRequests,
+		FailedRequests:  metricsResults.FailedRequests,
+		TotalDuration:   duration,
+		RequestsPerSec:  float64(metricsResults.TotalRequests) / duration.Seconds(),
+		Metrics:         metricsResults,
+	}
+
+	e.printResults(result)
+	return result, nil
+}
+
+// feedRequests sends requests to the worker pool
+func (e *Engine) feedRequests(requestChan chan<- struct{}) {
+	defer close(requestChan)
+
+	if e.config.Load.Requests > 0 {
+		// Fixed number of requests
+		for i := 0; i < e.config.Load.Requests; i++ {
+			select {
+			case requestChan <- struct{}{}:
+			case <-e.ctx.Done():
+				return
+			}
+		}
+	} else {
+		// Duration-based requests
+		ticker := time.NewTicker(time.Duration(1000/e.config.Load.RequestsPerSecond) * time.Millisecond)
+		defer ticker.Stop()
+
+		for {
+			select {
+			case <-ticker.C:
+				select {
+				case requestChan <- struct{}{}:
+				case <-e.ctx.Done():
+					return
+				}
+			case <-e.ctx.Done():
+				return
+			}
+		}
+	}
+}
+
+// run executes the worker loop
+func (w *Worker) run(requestChan <-chan struct{}) {
+	defer w.wg.Done()
+
+	for {
+		select {
+		case <-requestChan:
+			w.executeRequest()
+		case <-w.engine.ctx.Done():
+			return
+		}
+	}
+}
+
+// executeRequest performs a single HTTP request
+func (w *Worker) executeRequest() {
+	startTime := time.Now()
+
+	// Create request
+	var body io.Reader
+	if w.engine.config.Target.Body != "" {
+		body = strings.NewReader(w.engine.config.Target.Body)
+	}
+
+	req, err := http.NewRequestWithContext(
+		w.engine.ctx,
+		w.engine.config.Target.Method,
+		w.engine.config.Target.URL,
+		body,
+	)
+	if err != nil {
+		w.engine.metrics.RecordRequest(time.Since(startTime), 0, err)
+		return
+	}
+
+	// Add headers
+	for key, value := range w.engine.config.Target.Headers {
+		req.Header.Set(key, value)
+	}
+
+	// Execute request
+	resp, err := w.engine.client.Do(req)
+	duration := time.Since(startTime)
+
+	if err != nil {
+		w.engine.metrics.RecordRequest(duration, 0, err)
+		return
+	}
+	defer resp.Body.Close()
+
+	// Read response body (to ensure proper connection handling)
+	io.Copy(io.Discard, resp.Body)
+
+	// Record metrics
+	w.engine.metrics.RecordRequest(duration, resp.StatusCode, nil)
+}
+
+// Stop gracefully stops the engine
+func (e *Engine) Stop() {
+	e.cancel()
+}
+
+// printResults displays the test results
+func (e *Engine) printResults(result *Result) {
+	fmt.Printf("\n📊 Test Results:\n")
+	fmt.Printf("Duration: %.2fs | RPS: %.2f | Total: %d | Success: %d | Failed: %d\n",
+		result.TotalDuration.Seconds(),
+		result.RequestsPerSec,
+		result.TotalRequests,
+		result.SuccessRequests,
+		result.FailedRequests,
+	)
+
+	fmt.Printf("\nResponse Times:\n")
+	fmt.Printf("  Min: %v | Max: %v | Avg: %v\n",
+		result.Metrics.MinResponseTime,
+		result.Metrics.MaxResponseTime,
+		result.Metrics.AvgResponseTime,
+	)
+
+	fmt.Printf("  p50: %v | p90: %v | p95: %v | p99: %v\n",
+		result.Metrics.P50,
+		result.Metrics.P90,
+		result.Metrics.P95,
+		result.Metrics.P99,
+	)
+
+	if result.FailedRequests == 0 {
+		fmt.Printf("\n🎉 Fuck yeah! Your API handled it like a champ! 💪\n")
+	} else {
+		fmt.Printf("\n⚠️  Your API had some hiccups. Time to optimize! 🔧\n")
+	}
+}