Skip to content

Commit

Permalink
feat: preserve custom query parameters in influxdb_v2 output
Browse files Browse the repository at this point in the history
This also speeds up the write path by building the non-changing parts of
the "writeURL" once while still allowing for tag (or otherwise) defined
destination buckets.
  • Loading branch information
glinton committed Jun 7, 2024
1 parent 35edd18 commit a6b7936
Show file tree
Hide file tree
Showing 2 changed files with 94 additions and 17 deletions.
33 changes: 20 additions & 13 deletions plugins/outputs/influxdb_v2/http.go
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,7 @@ type httpClient struct {
client *http.Client
serializer *influx.Serializer
url *url.URL
params url.Values
retryTime time.Time
retryCount int
log telegraf.Logger
Expand Down Expand Up @@ -137,13 +138,19 @@ func NewHTTPClient(cfg *HTTPConfig) (*httpClient, error) {
return nil, fmt.Errorf("unsupported scheme %q", cfg.URL.Scheme)
}

preppedURL, params, err := prepWriteURL(*cfg.URL, cfg.Organization)
if err != nil {
return nil, err
}

client := &httpClient{
serializer: serializer,
client: &http.Client{
Timeout: timeout,
Transport: transport,
},
url: cfg.URL,
url: preppedURL,
params: params,
ContentEncoding: cfg.ContentEncoding,
Timeout: timeout,
Headers: headers,
Expand Down Expand Up @@ -246,15 +253,10 @@ func (c *httpClient) splitAndWriteBatch(ctx context.Context, bucket string, metr
}

func (c *httpClient) writeBatch(ctx context.Context, bucket string, metrics []telegraf.Metric) error {
loc, err := makeWriteURL(*c.url, c.Organization, bucket)
if err != nil {
return err
}

reader := c.requestBodyReader(metrics)
defer reader.Close()

req, err := c.makeWriteRequest(loc, reader)
req, err := c.makeWriteRequest(makeWriteURL(*c.url, c.params, bucket), reader)
if err != nil {
return err
}
Expand Down Expand Up @@ -402,11 +404,13 @@ func (c *httpClient) addHeaders(req *http.Request) {
}
}

func makeWriteURL(loc url.URL, org, bucket string) (string, error) {
params := url.Values{}
func makeWriteURL(loc url.URL, params url.Values, bucket string) string {
params.Set("bucket", bucket)
params.Set("org", org)
loc.RawQuery = params.Encode()
return loc.String()
}

func prepWriteURL(loc url.URL, org string) (*url.URL, url.Values, error) {
switch loc.Scheme {
case "unix":
loc.Scheme = "http"
Expand All @@ -415,10 +419,13 @@ func makeWriteURL(loc url.URL, org, bucket string) (string, error) {
case "http", "https":
loc.Path = path.Join(loc.Path, "/api/v2/write")
default:
return "", fmt.Errorf("unsupported scheme: %q", loc.Scheme)
return nil, nil, fmt.Errorf("unsupported scheme: %q", loc.Scheme)
}
loc.RawQuery = params.Encode()
return loc.String(), nil

params := loc.Query()
params.Set("org", org)

return &loc, params, nil
}

func (c *httpClient) Close() {
Expand Down
78 changes: 74 additions & 4 deletions plugins/outputs/influxdb_v2/http_internal_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ import (
"fmt"
"net/http"
"net/url"
"path"
"testing"
"time"

Expand All @@ -20,14 +21,26 @@ func TestMakeWriteURL(t *testing.T) {
err bool
url *url.URL
act string
bkt string
org string
}{
{
url: genURL("http://localhost:9999"),
act: "http://localhost:9999/api/v2/write?bucket=telegraf&org=influx",
act: "http://localhost:9999/api/v2/write?bucket=telegraf0&org=influx0",
bkt: "telegraf0",
org: "influx0",
},
{
url: genURL("http://localhost:9999?id=abc"),
act: "http://localhost:9999/api/v2/write?bucket=telegraf1&id=abc&org=influx1",
bkt: "telegraf1",
org: "influx1",
},
{
url: genURL("unix://var/run/influxd.sock"),
act: "http://127.0.0.1/api/v2/write?bucket=telegraf&org=influx",
act: "http://127.0.0.1/api/v2/write?bucket=telegraf2&org=influx2",
bkt: "telegraf2",
org: "influx2",
},
{
err: true,
Expand All @@ -36,15 +49,17 @@ func TestMakeWriteURL(t *testing.T) {
}

for i := range tests {
rURL, err := makeWriteURL(*tests[i].url, "influx", "telegraf")
rURL, params, err := prepWriteURL(*tests[i].url, tests[i].org)
if !tests[i].err {
require.NoError(t, err)
} else {
require.Error(t, err)
t.Log(err)
}
if err == nil {
require.Equal(t, tests[i].act, rURL)
for j := 0; j < 2; j++ {
require.Equal(t, tests[i].act, makeWriteURL(*rURL, params, tests[i].bkt))
}
}
}
}
Expand Down Expand Up @@ -98,3 +113,58 @@ func TestExponentialBackoffCalculationWithRetryAfter(t *testing.T) {
})
}
}

var (
bucket = "bkt"
org = "org"
loc, params, _ = prepWriteURL(*genURL("http://localhost:8086"), org)
)

// goos: linux
// goarch: amd64
// pkg: github.com/influxdata/telegraf/plugins/outputs/influxdb_v2
// cpu: 11th Gen Intel(R) Core(TM) i7-11850H @ 2.50GHz
// BenchmarkOldMakeWriteURL
// BenchmarkOldMakeWriteURL-16 1556631 683.2 ns/op 424 B/op 14 allocs/op
// PASS
// ok github.com/influxdata/telegraf/plugins/outputs/influxdb_v2 1.851s
func BenchmarkOldMakeWriteURL(b *testing.B) {
b.ReportAllocs()
for n := 0; n < b.N; n++ {
oldMakeWriteURL(*loc, org, bucket)
}
}

// goos: linux
// goarch: amd64
// pkg: github.com/influxdata/telegraf/plugins/outputs/influxdb_v2
// cpu: 11th Gen Intel(R) Core(TM) i7-11850H @ 2.50GHz
// BenchmarkNewMakeWriteURL
// BenchmarkNewMakeWriteURL-16 2084415 496.5 ns/op 280 B/op 9 allocs/op
// PASS
// ok github.com/influxdata/telegraf/plugins/outputs/influxdb_v2 1.626s
func BenchmarkNewMakeWriteURL(b *testing.B) {
b.ReportAllocs()
for n := 0; n < b.N; n++ {
makeWriteURL(*loc, params, bucket)
}
}

func oldMakeWriteURL(loc url.URL, org, bucket string) (string, error) {
params := url.Values{}
params.Set("bucket", bucket)
params.Set("org", org)

switch loc.Scheme {
case "unix":
loc.Scheme = "http"
loc.Host = "127.0.0.1"
loc.Path = "/api/v2/write"
case "http", "https":
loc.Path = path.Join(loc.Path, "/api/v2/write")
default:
return "", fmt.Errorf("unsupported scheme: %q", loc.Scheme)
}
loc.RawQuery = params.Encode()
return loc.String(), nil
}

0 comments on commit a6b7936

Please sign in to comment.