Cancel Http safe client requests before getting connection (#254)

* Fix race condition on gauge shutdown

* Refuse requests when trying to get connection instead of when it already
has it

Author: paulo

Makefile

@@ -15,9 +15,9 @@ test:
 integration-test:
 	docker-compose -f Dockercompose.test.yml up --build --abort-on-container-exit --always-recreate-deps
 	docker-compose -f Dockercompose.test.yml down --volumes
-	
+
 clean:
 	docker-compose -f Dockercompose.test.yml rm -f
 
 kafkacat:
-	docker run -it --network=host confluentinc/cp-kafkacat kafkacat -b localhost:19092 -C -t gotest -J
+	docker run -it --network=host confluentinc/cp-kafkacat kafkacat -b localhost:19092 -C -t gotest -J

http/http.go

@@ -52,9 +52,8 @@ func (no noLocalTransport) RoundTrip(req *http.Request) (*http.Response, error)
 	ctx, cancel := context.WithCancel(req.Context())
 
 	ctx = httptrace.WithClientTrace(ctx, &httptrace.ClientTrace{
-		GotConn: func(info httptrace.GotConnInfo) {
-			addr := info.Conn.RemoteAddr().String()
-			host, _, err := net.SplitHostPort(addr)
+		GetConn: func(hostPort string) {
+			host, _, err := net.SplitHostPort(hostPort)
 			if err != nil {
 				cancel()
 				no.errlog.WithError(err).Error("Cancelled request due to error in address parsing")

http/http_test.go

@@ -9,6 +9,7 @@ import (
 
 	"github.com/sirupsen/logrus"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func TestIsPrivateIP(t *testing.T) {
@@ -44,15 +45,16 @@ func TestSafeHTTPClient(t *testing.T) {
 
 	// It blocks the local IP.
 	_, err = client.Get(ts.URL)
-	assert.NotNil(t, err)
+	require.NotNil(t, err)
 
 	// It blocks localhost.
 	_, err = client.Get("http://localhost:" + tsURL.Port())
-	assert.NotNil(t, err)
+	require.NotNil(t, err)
 
 	// It works when reusing pooled connections.
 	for i := 0; i < 50; i++ {
-		_, err = client.Get("http://localhost:" + tsURL.Port())
+		res, err := client.Get("http://localhost:" + tsURL.Port())
+		assert.Nil(t, res)
 		assert.NotNil(t, err)
 	}
 

metriks/gauge.go

@@ -2,6 +2,7 @@ package metriks
 
 import (
 	"context"
+	"sync"
 	"sync/atomic"
 	"time"
 
@@ -25,6 +26,7 @@ type PersistentGauge struct {
 	ticker *time.Ticker
 	cancel context.CancelFunc
 	dur    time.Duration
+	wg     sync.WaitGroup
 }
 
 // Set will replace the value with a new one, it returns the old value
@@ -55,9 +57,11 @@ func (g *PersistentGauge) report(v int32) {
 }
 
 func (g *PersistentGauge) start(ctx context.Context) {
+	g.wg.Add(1)
 	for {
 		select {
 		case <-ctx.Done():
+			g.wg.Done()
 			return
 		case <-g.ticker.C:
 			g.report(g.value)
@@ -69,6 +73,7 @@ func (g *PersistentGauge) start(ctx context.Context) {
 // to the metrics collector
 func (g *PersistentGauge) Stop() {
 	g.cancel()
+	g.wg.Wait()
 }
 
 // NewPersistentGauge will create and start a PersistentGauge that reports the current value every 10s