Example of custom Go Agent Config.Transport to record agent network traffic

main.go

package main

import (
	"fmt"
	"io"
	"net/http"
	"os"
	"sync"
	"time"

	"github.com/newrelic/go-agent"
)

func mustGetEnv(key string) string {
	if val := os.Getenv(key); "" != val {
		return val
	}
	panic(fmt.Sprintf("environment variable %s unset", key))
}

type agentNetworkStats struct {
	sync.Mutex
	totalBytesSent int64
}

func (stats *agentNetworkStats) addBytesSent(n int64) {
	stats.Lock()
	defer stats.Unlock()
	stats.totalBytesSent += n
}

func (stats *agentNetworkStats) getTotalBytesSent() int64 {
	stats.Lock()
	defer stats.Unlock()
	return stats.totalBytesSent
}

type readerRecorder struct {
	r     io.ReadCloser
	stats *agentNetworkStats
}

func (r *readerRecorder) Close() error { return r.r.Close() }

func (r *readerRecorder) Read(p []byte) (int, error) {
	n, err := r.r.Read(p)
	r.stats.addBytesSent(int64(n))
	return n, err
}

type roundTripperFunc func(*http.Request) (*http.Response, error)

func (fn roundTripperFunc) RoundTrip(req *http.Request) (*http.Response, error) {
	return fn(req)
}

func main() {
	stats := &agentNetworkStats{}

	cfg := newrelic.NewConfig("Custom Config.Transport App", mustGetEnv("NEW_RELIC_LICENSE_KEY"))
	cfg.Logger = newrelic.NewDebugLogger(os.Stdout)

	cfg.Transport = roundTripperFunc(func(req *http.Request) (*http.Response, error) {
		if _, ok := req.Body.(*readerRecorder); !ok {
			req.Body = &readerRecorder{r: req.Body, stats: stats}
		}
		return http.DefaultTransport.RoundTrip(req)
	})

	app, err := newrelic.NewApplication(cfg)
	if nil != err {
		fmt.Println(err)
		os.Exit(1)
	}

	if err := app.WaitForConnection(5 * time.Second); nil != err {
		fmt.Println(err)
	}

	go func() {
		for {
			time.Sleep(30 * time.Second)
			sent := stats.getTotalBytesSent()
			fmt.Println("total agent bytes sent:", sent)
		}
	}()

	for {
		time.Sleep(5 * time.Second)
		app.RecordCustomEvent("task", map[string]interface{}{
			"color": "blue,",
		})
	}
}