monlet/agent/internal/client/client.go

package client

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"strings"
	"time"

	"github.com/monlet/agent/internal/scheduler"
)

const (
	MaxBatchEvents = 200
	MaxBatchBytes  = 1 << 20
)

type Client struct {
	baseURL string
	token   string
	agentID string
	http    *http.Client
}

func New(baseURL, token, agentID string) *Client {
	return &Client{
		baseURL: strings.TrimRight(baseURL, "/"),
		token:   token,
		agentID: agentID,
		http: &http.Client{
			Timeout:   15 * time.Second,
			Transport: &http.Transport{DisableKeepAlives: true},
		},
	}
}

type AgentFeatures struct {
	Push    bool `json:"push"`
	Metrics bool `json:"metrics"`
}

type HeartbeatRequest struct {
	AgentID    string            `json:"agent_id"`
	ObservedAt time.Time         `json:"observed_at"`
	Hostname   string            `json:"hostname"`
	Features   AgentFeatures     `json:"features"`
	Labels     map[string]string `json:"labels,omitempty"`
}

type EventBatchRequest struct {
	AgentID string            `json:"agent_id"`
	Events  []scheduler.Event `json:"events"`
}

type EventBatchResponse struct {
	Accepted     int `json:"accepted"`
	Deduplicated int `json:"deduplicated"`
}

func (c *Client) SendHeartbeat(ctx context.Context, hb HeartbeatRequest) error {
	return c.post(ctx, "/api/v1/heartbeat", hb, nil)
}

func (c *Client) SendEvents(ctx context.Context, events []scheduler.Event) (*EventBatchResponse, error) {
	if len(events) == 0 {
		return &EventBatchResponse{}, nil
	}
	if len(events) > MaxBatchEvents {
		return nil, nonRetryablef("batch exceeds %d events", MaxBatchEvents)
	}
	wireEvents := make([]scheduler.Event, len(events))
	for i, ev := range events {
		wireEvents[i] = ev.WireEvent()
	}
	req := EventBatchRequest{AgentID: c.agentID, Events: wireEvents}
	batches, err := c.splitEventBatches(req.Events)
	if err != nil {
		return nil, err
	}
	var total EventBatchResponse
	for _, batch := range batches {
		var resp EventBatchResponse
		if err := c.post(ctx, "/api/v1/events", EventBatchRequest{AgentID: c.agentID, Events: batch}, &resp); err != nil {
			return nil, err
		}
		total.Accepted += resp.Accepted
		total.Deduplicated += resp.Deduplicated
	}
	return &total, nil
}

func (c *Client) splitEventBatches(events []scheduler.Event) ([][]scheduler.Event, error) {
	overhead, err := c.eventBatchOverhead()
	if err != nil {
		return nil, err
	}
	batches := make([][]scheduler.Event, 0, 1)
	current := make([]scheduler.Event, 0, len(events))
	currentSize := overhead
	for _, ev := range events {
		eventSize, err := jsonSize(ev)
		if err != nil {
			return nil, err
		}
		nextSize := currentSize + eventSize
		if len(current) > 0 {
			nextSize++
		}
		if nextSize <= MaxBatchBytes {
			current = append(current, ev)
			currentSize = nextSize
			continue
		}
		if len(current) == 0 {
			return nil, nonRetryablef("single event batch exceeds %d bytes", MaxBatchBytes)
		}
		batches = append(batches, current)
		current = []scheduler.Event{ev}
		currentSize = overhead + eventSize
		if currentSize > MaxBatchBytes {
			return nil, nonRetryablef("single event batch exceeds %d bytes", MaxBatchBytes)
		}
	}
	if len(current) > 0 {
		batches = append(batches, current)
	}
	return batches, nil
}

func (c *Client) eventBatchOverhead() (int, error) {
	size, err := jsonSize(EventBatchRequest{AgentID: c.agentID, Events: []scheduler.Event{}})
	if err != nil {
		return 0, err
	}
	return size, nil
}

func jsonSize(v any) (int, error) {
	b, err := json.Marshal(v)
	if err != nil {
		return 0, err
	}
	return len(b), nil
}

func (c *Client) post(ctx context.Context, path string, body any, out any) error {
	b, err := json.Marshal(body)
	if err != nil {
		return err
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL+path, bytes.NewReader(b))
	if err != nil {
		return err
	}
	req.Header.Set("Authorization", "Bearer "+c.token)
	req.Header.Set("Content-Type", "application/json")
	resp, err := c.http.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode/100 != 2 {
		msg, _ := io.ReadAll(io.LimitReader(resp.Body, 1024))
		return &HTTPError{Status: resp.StatusCode, Body: string(msg)}
	}
	if out != nil {
		return json.NewDecoder(resp.Body).Decode(out)
	}
	return nil
}

type HTTPError struct {
	Status int
	Body   string
}

func (e *HTTPError) Error() string { return fmt.Sprintf("http %d: %s", e.Status, e.Body) }

type NonRetryableError struct{ Err error }

func (e *NonRetryableError) Error() string { return e.Err.Error() }
func (e *NonRetryableError) Unwrap() error { return e.Err }

func nonRetryablef(format string, args ...any) error {
	return &NonRetryableError{Err: fmt.Errorf(format, args...)}
}

// Retryable reports whether the error should trigger a retry (network errors
// and 5xx are retryable; 4xx are dropped to avoid hot-loop on bad payload).
func Retryable(err error) bool {
	if err == nil {
		return false
	}
	var ne *NonRetryableError
	if errors.As(err, &ne) {
		return false
	}
	var he *HTTPError
	if errors.As(err, &he) {
		return he.Status >= 500
	}
	return true
}