dom
/
telegraf
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

Fix influxdb output serialization on connection closed (#6621)

This commit is contained in:
陈方舟 2019-11-14 04:56:01 +08:00 committed by Daniel Nelson
parent 9a2b3bc917
commit fa2f0fff4e
8 changed files with 132 additions and 37 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
internal/internal.go
View File

@ -16,6 +16,7 @@ import (
"runtime"
"strconv"
"strings"
"sync"
"syscall"
"time"
"unicode"
@ -50,6 +51,11 @@ type Number struct {
Value float64
}
type ReadWaitCloser struct {
pipeReader *io.PipeReader
wg sync.WaitGroup
}
// SetVersion sets the telegraf agent version
func SetVersion(v string) error {
if version != "" {
@ -281,14 +287,25 @@ func ExitStatus(err error) (int, bool) {
return 0, false
}
func (r *ReadWaitCloser) Close() error {
err := r.pipeReader.Close()
r.wg.Wait() // wait for the gzip goroutine finish
return err
}
// CompressWithGzip takes an io.Reader as input and pipes
// it through a gzip.Writer returning an io.Reader containing
// the gzipped data.
// An error is returned if passing data to the gzip.Writer fails
func CompressWithGzip(data io.Reader) (io.Reader, error) {
func CompressWithGzip(data io.Reader) (io.ReadCloser, error) {
pipeReader, pipeWriter := io.Pipe()
gzipWriter := gzip.NewWriter(pipeWriter)
rc := &ReadWaitCloser{
pipeReader: pipeReader,
}
rc.wg.Add(1)
var err error
go func() {
_, err = io.Copy(gzipWriter, data)
@ -296,6 +313,7 @@ func CompressWithGzip(data io.Reader) (io.Reader, error) {
// subsequent reads from the read half of the pipe will
// return no bytes and the error err, or EOF if err is nil.
pipeWriter.CloseWithError(err)
rc.wg.Done()
}()
return pipeReader, err

34
internal/internal_test.go
View File

@ -3,6 +3,8 @@ package internal
import (
"bytes"
"compress/gzip"
"crypto/rand"
"io"
"io/ioutil"
"log"
"os/exec"
@ -232,6 +234,38 @@ func TestCompressWithGzip(t *testing.T) {
assert.Equal(t, testData, string(output))
}
type mockReader struct {
readN uint64 // record the number of calls to Read
}
func (r *mockReader) Read(p []byte) (n int, err error) {
r.readN++
return rand.Read(p)
}
func TestCompressWithGzipEarlyClose(t *testing.T) {
mr := &mockReader{}
rc, err := CompressWithGzip(mr)
assert.NoError(t, err)
n, err := io.CopyN(ioutil.Discard, rc, 10000)
assert.NoError(t, err)
assert.Equal(t, int64(10000), n)
r1 := mr.readN
err = rc.Close()
assert.NoError(t, err)
n, err = io.CopyN(ioutil.Discard, rc, 10000)
assert.Error(t, io.EOF, err)
assert.Equal(t, int64(0), n)
r2 := mr.readN
// no more read to the source after closing
assert.Equal(t, r1, r2)
}
func TestVersionAlreadySet(t *testing.T) {
err := SetVersion("foo")
assert.Nil(t, err)

9
plugins/inputs/http/http.go
View File

@ -153,6 +153,7 @@ func (h *HTTP) gatherURL(
if err != nil {
return err
}
defer body.Close()
request, err := http.NewRequest(h.Method, url, body)
if err != nil {
@ -216,16 +217,16 @@ func (h *HTTP) gatherURL(
return nil
}
func makeRequestBodyReader(contentEncoding, body string) (io.Reader, error) {
var err error
func makeRequestBodyReader(contentEncoding, body string) (io.ReadCloser, error) {
var reader io.Reader = strings.NewReader(body)
if contentEncoding == "gzip" {
reader, err = internal.CompressWithGzip(reader)
rc, err := internal.CompressWithGzip(reader)
if err != nil {
return nil, err
}
return rc, nil
}
return reader, nil
return ioutil.NopCloser(reader), nil
}
func init() {

4
plugins/outputs/http/http.go
View File

@ -176,10 +176,12 @@ func (h *HTTP) write(reqBody []byte) error {
var err error
if h.ContentEncoding == "gzip" {
reqBodyBuffer, err = internal.CompressWithGzip(reqBodyBuffer)
rc, err := internal.CompressWithGzip(reqBodyBuffer)
if err != nil {
return err
}
defer rc.Close()
reqBodyBuffer = rc
}
req, err := http.NewRequest(h.Method, h.URL, reqBodyBuffer)

31
plugins/outputs/influxdb/http.go
View File

@ -6,6 +6,7 @@ import (
"encoding/json"
"fmt"
"io"
"io/ioutil"
"net"
"net/http"
"net/url"
@ -288,7 +289,12 @@ func (c *httpClient) writeBatch(ctx context.Context, db string, metrics []telegr
return err
}
reader := influx.NewReader(metrics, c.config.Serializer)
reader, err := c.requestBodyReader(metrics)
if err != nil {
return err
}
defer reader.Close()
req, err := c.makeWriteRequest(url, reader)
if err != nil {
return err
@ -386,12 +392,6 @@ func (c *httpClient) makeQueryRequest(query string) (*http.Request, error) {
func (c *httpClient) makeWriteRequest(url string, body io.Reader) (*http.Request, error) {
var err error
if c.config.ContentEncoding == "gzip" {
body, err = internal.CompressWithGzip(body)
if err != nil {
return nil, err
}
}
req, err := http.NewRequest("POST", url, body)
if err != nil {
@ -408,6 +408,23 @@ func (c *httpClient) makeWriteRequest(url string, body io.Reader) (*http.Request
return req, nil
}
// requestBodyReader warp io.Reader from influx.NewReader to io.ReadCloser, which is usefully to fast close the write
// side of the connection in case of error
func (c *httpClient) requestBodyReader(metrics []telegraf.Metric) (io.ReadCloser, error) {
reader := influx.NewReader(metrics, c.config.Serializer)
if c.config.ContentEncoding == "gzip" {
rc, err := internal.CompressWithGzip(reader)
if err != nil {
return nil, err
}
return rc, nil
}
return ioutil.NopCloser(reader), nil
}
func (c *httpClient) addHeaders(req *http.Request) {
if c.config.Username != "" || c.config.Password != "" {
req.SetBasicAuth(c.config.Username, c.config.Password)

19
plugins/outputs/influxdb/influxdb.go
View File

@ -57,7 +57,6 @@ type InfluxDB struct {
CreateHTTPClientF func(config *HTTPConfig) (Client, error)
CreateUDPClientF func(config *UDPConfig) (Client, error)
serializer *influx.Serializer
Log telegraf.Logger
}
@ -145,11 +144,6 @@ func (i *InfluxDB) Connect() error {
urls = append(urls, defaultURL)
}
i.serializer = influx.NewSerializer()
if i.InfluxUintSupport {
i.serializer.SetFieldTypeSupport(influx.UintSupport)
}
for _, u := range urls {
parts, err := url.Parse(u)
if err != nil {
@ -237,7 +231,7 @@ func (i *InfluxDB) udpClient(url *url.URL) (Client, error) {
config := &UDPConfig{
URL: url,
MaxPayloadSize: int(i.UDPPayload.Size),
Serializer: i.serializer,
Serializer: i.newSerializer(),
Log: i.Log,
}
@ -271,7 +265,7 @@ func (i *InfluxDB) httpClient(ctx context.Context, url *url.URL, proxy *url.URL)
SkipDatabaseCreation: i.SkipDatabaseCreation,
RetentionPolicy: i.RetentionPolicy,
Consistency: i.WriteConsistency,
Serializer: i.serializer,
Serializer: i.newSerializer(),
Log: i.Log,
}
@ -291,6 +285,15 @@ func (i *InfluxDB) httpClient(ctx context.Context, url *url.URL, proxy *url.URL)
return c, nil
}
func (i *InfluxDB) newSerializer() *influx.Serializer {
serializer := influx.NewSerializer()
if i.InfluxUintSupport {
serializer.SetFieldTypeSupport(influx.UintSupport)
}
return serializer
}
func init() {
outputs.Add("influxdb", func() telegraf.Output {
return &InfluxDB{

31
plugins/outputs/influxdb_v2/http.go
View File

@ -7,6 +7,7 @@ import (
"errors"
"fmt"
"io"
"io/ioutil"
"log"
"net"
"net/http"
@ -214,7 +215,12 @@ func (c *httpClient) writeBatch(ctx context.Context, bucket string, metrics []te
return err
}
reader := influx.NewReader(metrics, c.serializer)
reader, err := c.requestBodyReader(metrics)
if err != nil {
return err
}
defer reader.Close()
req, err := c.makeWriteRequest(url, reader)
if err != nil {
return err
@ -282,12 +288,6 @@ func (c *httpClient) writeBatch(ctx context.Context, bucket string, metrics []te
func (c *httpClient) makeWriteRequest(url string, body io.Reader) (*http.Request, error) {
var err error
if c.ContentEncoding == "gzip" {
body, err = internal.CompressWithGzip(body)
if err != nil {
return nil, err
}
}
req, err := http.NewRequest("POST", url, body)
if err != nil {
@ -304,6 +304,23 @@ func (c *httpClient) makeWriteRequest(url string, body io.Reader) (*http.Request
return req, nil
}
// requestBodyReader warp io.Reader from influx.NewReader to io.ReadCloser, which is usefully to fast close the write
// side of the connection in case of error
func (c *httpClient) requestBodyReader(metrics []telegraf.Metric) (io.ReadCloser, error) {
reader := influx.NewReader(metrics, c.serializer)
if c.ContentEncoding == "gzip" {
rc, err := internal.CompressWithGzip(reader)
if err != nil {
return nil, err
}
return rc, nil
}
return ioutil.NopCloser(reader), nil
}
func (c *httpClient) addHeaders(req *http.Request) {
for header, value := range c.Headers {
req.Header.Set(header, value)

17
plugins/outputs/influxdb_v2/influxdb.go
View File

@ -97,7 +97,6 @@ type InfluxDB struct {
tls.ClientConfig
clients []Client
serializer *influx.Serializer
}
func (i *InfluxDB) Connect() error {
@ -107,11 +106,6 @@ func (i *InfluxDB) Connect() error {
i.URLs = append(i.URLs, defaultURL)
}
i.serializer = influx.NewSerializer()
if i.UintSupport {
i.serializer.SetFieldTypeSupport(influx.UintSupport)
}
for _, u := range i.URLs {
parts, err := url.Parse(u)
if err != nil {
@ -196,7 +190,7 @@ func (i *InfluxDB) getHTTPClient(ctx context.Context, url *url.URL, proxy *url.U
UserAgent: i.UserAgent,
ContentEncoding: i.ContentEncoding,
TLSConfig: tlsConfig,
Serializer: i.serializer,
Serializer: i.newSerializer(),
}
c, err := NewHTTPClient(config)
@ -207,6 +201,15 @@ func (i *InfluxDB) getHTTPClient(ctx context.Context, url *url.URL, proxy *url.U
return c, nil
}
func (i *InfluxDB) newSerializer() *influx.Serializer {
serializer := influx.NewSerializer()
if i.UintSupport {
serializer.SetFieldTypeSupport(influx.UintSupport)
}
return serializer
}
func init() {
outputs.Add("influxdb_v2", func() telegraf.Output {
return &InfluxDB{