115 lines
2.8 KiB
Go
115 lines
2.8 KiB
Go
package main
|
|
|
|
import (
|
|
"bufio"
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
type namedURLs []string
|
|
|
|
func (n *namedURLs) String() string { return strings.Join(*n, ",") }
|
|
func (n *namedURLs) Set(value string) error {
|
|
if !strings.Contains(value, "=") {
|
|
return fmt.Errorf("URL must use name=https://host/path")
|
|
}
|
|
*n = append(*n, value)
|
|
return nil
|
|
}
|
|
|
|
type metricSnapshot struct {
|
|
Timestamp time.Time `json:"timestamp"`
|
|
Name string `json:"name"`
|
|
URL string `json:"url"`
|
|
Status int `json:"status"`
|
|
Duration int64 `json:"duration_ms"`
|
|
Body string `json:"body,omitempty"`
|
|
Error string `json:"error,omitempty"`
|
|
}
|
|
|
|
func runCollector(ctx context.Context, sources namedURLs, output string, interval, duration time.Duration) error {
|
|
if len(sources) == 0 {
|
|
return fmt.Errorf("at least one --url name=url is required")
|
|
}
|
|
if interval <= 0 || duration <= 0 {
|
|
return fmt.Errorf("interval and duration must be positive")
|
|
}
|
|
if dir := filepath.Dir(output); dir != "." {
|
|
if err := os.MkdirAll(dir, 0o750); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
file, err := os.OpenFile(output, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0o600)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer file.Close()
|
|
writer := bufio.NewWriterSize(file, 256<<10)
|
|
defer writer.Flush()
|
|
client := &http.Client{Timeout: 10 * time.Second}
|
|
deadline := time.Now().Add(duration)
|
|
for {
|
|
snapshots := collectOnce(ctx, client, sources)
|
|
for _, snapshot := range snapshots {
|
|
raw, _ := json.Marshal(snapshot)
|
|
if _, err = writer.Write(append(raw, '\n')); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
if err = writer.Flush(); err != nil {
|
|
return err
|
|
}
|
|
if time.Now().Add(interval).After(deadline) {
|
|
return nil
|
|
}
|
|
if err = waitContext(ctx, interval); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
}
|
|
|
|
func collectOnce(ctx context.Context, client *http.Client, sources namedURLs) []metricSnapshot {
|
|
out := make([]metricSnapshot, len(sources))
|
|
var wg sync.WaitGroup
|
|
for i, source := range sources {
|
|
i, source := i, source
|
|
wg.Add(1)
|
|
go func() {
|
|
defer wg.Done()
|
|
parts := strings.SplitN(source, "=", 2)
|
|
snapshot := metricSnapshot{Timestamp: time.Now().UTC(), Name: parts[0], URL: parts[1]}
|
|
started := time.Now()
|
|
req, err := http.NewRequestWithContext(ctx, http.MethodGet, parts[1], nil)
|
|
if err == nil {
|
|
var response *http.Response
|
|
response, err = client.Do(req)
|
|
if err == nil {
|
|
snapshot.Status = response.StatusCode
|
|
raw, readErr := io.ReadAll(io.LimitReader(response.Body, 16<<20))
|
|
_ = response.Body.Close()
|
|
if readErr != nil {
|
|
err = readErr
|
|
} else {
|
|
snapshot.Body = string(raw)
|
|
}
|
|
}
|
|
}
|
|
snapshot.Duration = time.Since(started).Milliseconds()
|
|
if err != nil {
|
|
snapshot.Error = err.Error()
|
|
}
|
|
out[i] = snapshot
|
|
}()
|
|
}
|
|
wg.Wait()
|
|
return out
|
|
}
|