Logging-System/sink/httpbatch.go

package sink

import (
	"bytes"
	"encoding/json"
	"errors"
	"net/http"
	"sync"
	"sync/atomic"
	"time"

	"logging-system/logsys"
)

// HTTPBatch 把日志攒批 POST 到采集器。Async 为 false 时在 Write 里同步发送。
type HTTPBatch struct {
	URL           string
	Token         string
	MaxEntries    int
	FlushInterval time.Duration
	MaxRetries    int
	Async         bool
	Client        *http.Client
	FlushTimeout  time.Duration

	mu      sync.Mutex
	buf     []logsys.Entry
	stop    chan struct{}
	done    chan struct{}
	started bool
	dropped atomic.Int64
	once    sync.Once
}

// Start 启动异步刷盘。同步模式不必调用。
func (h *HTTPBatch) Start() {
	if !h.Async {
		h.defaults()
		return
	}
	h.once.Do(func() {
		h.defaults()
		h.stop = make(chan struct{})
		h.done = make(chan struct{})
		h.started = true
		go h.loop()
	})
}

func (h *HTTPBatch) defaults() {
	if h.MaxEntries < 1 {
		h.MaxEntries = 200
	}
	if h.FlushInterval <= 0 {
		h.FlushInterval = 2 * time.Second
	}
	if h.MaxRetries < 1 {
		h.MaxRetries = 3
	}
	if h.Client == nil {
		h.Client = &http.Client{Timeout: 5 * time.Second}
	}
	if h.FlushTimeout <= 0 {
		h.FlushTimeout = 5 * time.Second
	}
}

// Write 追加一条日志。同步模式立即发送；异步模式达到批量上限时触发发送。
func (h *HTTPBatch) Write(e logsys.Entry) error {
	h.Start()
	if !h.Async {
		return h.post([]logsys.Entry{e})
	}
	h.mu.Lock()
	h.buf = append(h.buf, e)
	full := len(h.buf) >= h.MaxEntries
	var batch []logsys.Entry
	if full {
		batch = h.take()
	}
	h.mu.Unlock()
	if full {
		if err := h.post(batch); err != nil {
			h.dropped.Add(int64(len(batch)))
			return err
		}
	}
	return nil
}

// Dropped 返回发送失败后丢弃的条数。
func (h *HTTPBatch) Dropped() int64 { return h.dropped.Load() }

// Sync 立刻送出缓冲。
func (h *HTTPBatch) Sync() error {
	h.Start()
	h.mu.Lock()
	batch := h.take()
	h.mu.Unlock()
	if len(batch) == 0 {
		return nil
	}
	if err := h.post(batch); err != nil {
		h.dropped.Add(int64(len(batch)))
		return err
	}
	return nil
}

// Close 停止后台循环并送出剩余日志。
func (h *HTTPBatch) Close() error {
	h.Start()
	if h.Async && h.started {
		select {
		case <-h.stop:
		default:
			close(h.stop)
		}
		select {
		case <-h.done:
		case <-time.After(h.FlushTimeout):
			return errors.New("sink: http batch close timeout")
		}
	}
	return h.Sync()
}

func (h *HTTPBatch) loop() {
	defer close(h.done)
	ticker := time.NewTicker(h.FlushInterval)
	defer ticker.Stop()
	for {
		select {
		case <-h.stop:
			return
		case <-ticker.C:
			_ = h.Sync()
		}
	}
}

func (h *HTTPBatch) take() []logsys.Entry {
	if len(h.buf) == 0 {
		return nil
	}
	out := h.buf
	h.buf = nil
	return out
}

func (h *HTTPBatch) post(batch []logsys.Entry) error {
	if len(batch) == 0 {
		return nil
	}
	if h.URL == "" {
		return errors.New("sink: http url is required")
	}
	body, err := json.Marshal(map[string]any{"entries": batch})
	if err != nil {
		return err
	}
	var last error
	for attempt := 0; attempt < h.MaxRetries; attempt++ {
		req, err := http.NewRequest(http.MethodPost, h.URL, bytes.NewReader(body))
		if err != nil {
			return err
		}
		req.Header.Set("Content-Type", "application/json")
		if h.Token != "" {
			req.Header.Set("Authorization", "Bearer "+h.Token)
		}
		res, err := h.Client.Do(req)
		if err == nil {
			_ = res.Body.Close()
			if res.StatusCode >= 200 && res.StatusCode < 300 {
				return nil
			}
			err = errors.New(res.Status)
		}
		last = err
		if attempt+1 < h.MaxRetries {
			time.Sleep(time.Duration(attempt+1) * 10 * time.Millisecond)
		}
	}
	return last
}