ginuerzh
64a16fafbc
- httpRecorder: use http.NewRequestWithContext instead of http.NewRequest - tcpRecorder: loop Write to avoid short-write data loss - fileRecorder: guard against nil out in Record and Close - plugin/grpc, plugin/http: handle json.Marshal errors instead of ignoring - fileRecorder, redis*Recorder: idempotent Close via sync.Once - All exported symbols: add doc comments
86 lines
1.9 KiB
Go
86 lines
1.9 KiB
Go
package recorder
|
|
|
|
import (
|
|
"context"
|
|
"net"
|
|
"time"
|
|
|
|
"github.com/go-gost/core/logger"
|
|
"github.com/go-gost/core/metrics"
|
|
"github.com/go-gost/core/recorder"
|
|
xmetrics "github.com/go-gost/x/metrics"
|
|
)
|
|
|
|
type tcpRecorderOptions struct {
|
|
recorder string
|
|
timeout time.Duration
|
|
log logger.Logger
|
|
}
|
|
|
|
// TCPRecorderOption configures TCPRecorder options.
|
|
type TCPRecorderOption func(opts *tcpRecorderOptions)
|
|
|
|
// RecorderTCPRecorderOption sets the recorder name for metrics labeling.
|
|
func RecorderTCPRecorderOption(recorder string) TCPRecorderOption {
|
|
return func(opts *tcpRecorderOptions) {
|
|
opts.recorder = recorder
|
|
}
|
|
}
|
|
|
|
// TimeoutTCPRecorderOption sets the dial timeout for TCP connections.
|
|
func TimeoutTCPRecorderOption(timeout time.Duration) TCPRecorderOption {
|
|
return func(opts *tcpRecorderOptions) {
|
|
opts.timeout = timeout
|
|
}
|
|
}
|
|
|
|
// LogTCPRecorderOption sets the logger for the TCP recorder.
|
|
func LogTCPRecorderOption(log logger.Logger) TCPRecorderOption {
|
|
return func(opts *tcpRecorderOptions) {
|
|
opts.log = log
|
|
}
|
|
}
|
|
|
|
type tcpRecorder struct {
|
|
recorder string
|
|
addr string
|
|
dialer *net.Dialer
|
|
log logger.Logger
|
|
}
|
|
|
|
// TCPRecorder records data to TCP service.
|
|
func TCPRecorder(addr string, opts ...TCPRecorderOption) recorder.Recorder {
|
|
var options tcpRecorderOptions
|
|
for _, opt := range opts {
|
|
opt(&options)
|
|
}
|
|
|
|
return &tcpRecorder{
|
|
recorder: options.recorder,
|
|
addr: addr,
|
|
dialer: &net.Dialer{
|
|
Timeout: options.timeout,
|
|
},
|
|
log: options.log,
|
|
}
|
|
}
|
|
|
|
func (r *tcpRecorder) Record(ctx context.Context, b []byte, opts ...recorder.RecordOption) error {
|
|
xmetrics.GetCounter(xmetrics.MetricRecorderRecordsCounter, metrics.Labels{"recorder": r.recorder}).Inc()
|
|
|
|
c, err := r.dialer.DialContext(ctx, "tcp", r.addr)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer c.Close()
|
|
|
|
for written := 0; written < len(b); {
|
|
n, err := c.Write(b[written:])
|
|
if err != nil {
|
|
return err
|
|
}
|
|
written += n
|
|
}
|
|
return nil
|
|
}
|