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
91 lines
1.8 KiB
Go
91 lines
1.8 KiB
Go
package recorder
|
|
|
|
import (
|
|
"context"
|
|
"io"
|
|
"sync"
|
|
|
|
"github.com/go-gost/core/metrics"
|
|
"github.com/go-gost/core/recorder"
|
|
xmetrics "github.com/go-gost/x/metrics"
|
|
)
|
|
|
|
type fileRecorderOptions struct {
|
|
recorder string
|
|
sep string
|
|
}
|
|
|
|
// FileRecorderOption configures FileRecorder options.
|
|
type FileRecorderOption func(opts *fileRecorderOptions)
|
|
|
|
// RecorderFileRecorderOption sets the recorder name for metrics labeling.
|
|
func RecorderFileRecorderOption(recorder string) FileRecorderOption {
|
|
return func(opts *fileRecorderOptions) {
|
|
opts.recorder = recorder
|
|
}
|
|
}
|
|
|
|
// SepFileRecorderOption sets the record separator written after each record.
|
|
func SepFileRecorderOption(sep string) FileRecorderOption {
|
|
return func(opts *fileRecorderOptions) {
|
|
opts.sep = sep
|
|
}
|
|
}
|
|
|
|
type fileRecorder struct {
|
|
recorder string
|
|
out io.WriteCloser
|
|
sep string
|
|
|
|
mu sync.Mutex
|
|
closeOnce sync.Once
|
|
}
|
|
|
|
// FileRecorder records data to file.
|
|
func FileRecorder(out io.WriteCloser, opts ...FileRecorderOption) recorder.Recorder {
|
|
var options fileRecorderOptions
|
|
for _, opt := range opts {
|
|
opt(&options)
|
|
}
|
|
|
|
return &fileRecorder{
|
|
recorder: options.recorder,
|
|
out: out,
|
|
sep: options.sep,
|
|
}
|
|
}
|
|
|
|
func (r *fileRecorder) Record(ctx context.Context, b []byte, opts ...recorder.RecordOption) error {
|
|
if r.out == nil {
|
|
return nil
|
|
}
|
|
|
|
xmetrics.GetCounter(xmetrics.MetricRecorderRecordsCounter, metrics.Labels{"recorder": r.recorder}).Inc()
|
|
|
|
if r.sep != "" {
|
|
r.mu.Lock() // mutex is used to prevent unordered writes to the file
|
|
defer r.mu.Unlock()
|
|
}
|
|
|
|
if _, err := r.out.Write(b); err != nil {
|
|
return err
|
|
}
|
|
|
|
if r.sep != "" {
|
|
_, err := io.WriteString(r.out, r.sep)
|
|
return err
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func (r *fileRecorder) Close() error {
|
|
var err error
|
|
r.closeOnce.Do(func() {
|
|
if r.out != nil {
|
|
err = r.out.Close()
|
|
}
|
|
})
|
|
return err
|
|
}
|