Optimize format and add more test

Signed-off-by: Filip Petkovski <filip.petkovsky@gmail.com>

pkg/receive/handler.go

@@ -406,6 +406,7 @@ func (h *Handler) getStats(r *http.Request, statsByLabelName string) ([]statusap
 
 // Close stops the Handler.
 func (h *Handler) Close() {
+	_ = h.peers.Close()
 	runutil.CloseWithLogOnErr(h.logger, h.httpSrv, "receive HTTP server")
 }
 
@@ -1330,7 +1331,7 @@ type peerWorker struct {
 	forwardDelay prometheus.Histogram
 }
 
-func newPeerGroup(backoff backoff.Backoff, forwardDelay prometheus.Histogram, asyncForwardWorkersCount uint, useCapNProtoReplication bool, dialOpts ...grpc.DialOption) peersContainer {
+func newPeerGroup(backoff backoff.Backoff, forwardDelay prometheus.Histogram, asyncForwardWorkersCount uint, useCapNProtoReplication bool, dialOpts ...grpc.DialOption) *peerGroup {
 	return &peerGroup{
 		dialOpts:                 dialOpts,
 		connections:              map[Endpoint]*peerWorker{},
@@ -1350,6 +1351,7 @@ type peersContainer interface {
 	markPeerUnavailable(Endpoint)
 	markPeerAvailable(Endpoint)
 	reset()
+	io.Closer
 }
 
 func (p *peerWorker) RemoteWriteAsync(ctx context.Context, req *storepb.WriteRequest, er endpointReplica, seriesIDs []int, responseWriter chan writeResponse, cb func(error)) {
@@ -1392,6 +1394,13 @@ type peerGroup struct {
 	dialer func(target string, opts ...grpc.DialOption) (conn *grpc.ClientConn, err error)
 }
 
+func (p *peerGroup) Close() error {
+	for _, c := range p.connections {
+		c.wp.Close()
+	}
+	return nil
+}
+
 func (p *peerGroup) close(endpoint Endpoint) error {
 	p.m.Lock()
 	defer p.m.Unlock()

pkg/receive/handler_test.go

@@ -177,6 +177,13 @@ type fakePeersGroup struct {
 	closeCalled map[Endpoint]bool
 }
 
+func (g *fakePeersGroup) Close() error {
+	for _, c := range g.clients {
+		c.wp.Close()
+	}
+	return nil
+}
+
 func (g *fakePeersGroup) markPeerUnavailable(s Endpoint) {
 }
 
@@ -253,19 +260,13 @@ func newTestHandlerHashring(
 			peer = newPeerWorker(client, prometheus.NewHistogram(prometheus.HistogramOpts{}), 1)
 			closers = append(closers, func() error {
 				srv.Shutdown()
-				peer.wp.Close()
 				return goerrors.Join(listener.Close(), client.Close())
 			})
 			go func() { _ = srv.ListenAndServe() }()
-			fakePeers.clients[endpoint] = peer
 		} else {
 			peer = newPeerWorker(&fakeRemoteWriteGRPCServer{h: h}, prometheus.NewHistogram(prometheus.HistogramOpts{}), 1)
-			closers = append(closers, func() error {
-				peer.wp.Close()
-				return nil
-			})
-			fakePeers.clients[endpoint] = peer
 		}
+		fakePeers.clients[endpoint] = peer
 	}
 	// Use hashmod as default.
 	if hashringAlgo == "" {
@@ -663,6 +664,12 @@ func testReceiveQuorum(t *testing.T, hashringAlgo HashringAlgorithm, withConsist
 			}
 			defer func() {
 				testutil.Ok(t, closeFunc())
+				// Wait a few milliseconds for peer workers to process the queue.
+				time.AfterFunc(50*time.Millisecond, func() {
+					for _, h := range handlers {
+						h.Close()
+					}
+				})
 			}()
 			tenant := "test"
 

pkg/receive/writecapnp/client.go

@@ -8,6 +8,7 @@ import (
 	"net"
 	"sync"
 
+	"capnproto.org/go/capnp/v3"
 	"capnproto.org/go/capnp/v3/rpc"
 	"github.com/pkg/errors"
 	"google.golang.org/grpc"
@@ -45,7 +46,7 @@ type RemoteWriteClient struct {
 	mu sync.Mutex
 
 	dialer Dialer
-	conn   *net.Conn
+	conn   *rpc.Conn
 
 	writer Writer
 }
@@ -64,19 +65,32 @@ func (r *RemoteWriteClient) writeWithReconnect(ctx context.Context, numReconnect
 	if err := r.connect(ctx); err != nil {
 		return nil, err
 	}
+	arena := capnp.SingleSegment(nil)
+	defer arena.Release()
+
 	result, release := r.writer.Write(ctx, func(params Writer_write_Params) error {
-		wr, err := Build(in.Tenant, in.Timeseries)
+		_, seg, err := capnp.NewMessage(arena)
+		if err != nil {
+			return err
+		}
+		wr, err := NewRootWriteRequest(seg)
 		if err != nil {
 			return err
 		}
-		return params.SetWr(wr)
+		if err := params.SetWr(wr); err != nil {
+			return err
+		}
+		wr, err = params.Wr()
+		if err != nil {
+			return err
+		}
+		return BuildInto(wr, in.Tenant, in.Timeseries)
 	})
 	defer release()
 
 	s, err := result.Struct()
 	if err != nil {
-		if numReconnects > 0 {
-			_ = r.Close()
+		if numReconnects > 0 && capnp.IsDisconnected(err) {
 			if err := r.Close(); err != nil {
 				return nil, err
 			}
@@ -110,20 +124,18 @@ func (r *RemoteWriteClient) connect(ctx context.Context) error {
 	if err != nil {
 		return errors.Wrap(err, "failed to dial peer")
 	}
-	r.writer = Writer(rpc.NewConn(rpc.NewPackedStreamTransport(conn), nil).Bootstrap(ctx))
-	r.conn = &conn
+	r.conn = rpc.NewConn(rpc.NewPackedStreamTransport(conn), nil)
+	r.writer = Writer(r.conn.Bootstrap(ctx))
 	return nil
 }
 
 func (r *RemoteWriteClient) Close() error {
 	r.mu.Lock()
-	defer r.mu.Unlock()
 	if r.conn != nil {
-		r.writer.Release()
-
-		conn := *r.conn
+		conn := r.conn
 		r.conn = nil
-		return conn.Close()
+		go conn.Close()
 	}
+	r.mu.Unlock()
 	return nil
 }