[ES-1122621] fix receive connection re-establish taking too long (5m) (#36)

Author: jnyi

cmd/thanos/receive.go

@@ -158,6 +158,7 @@ func runReceive(
 		dialOpts = append(dialOpts, grpc.WithDefaultCallOptions(grpc.UseCompressor(conf.compression)))
 	}
 	if receiveMode == receive.RouterOnly {
+		dialOpts = append(dialOpts, grpc.WithIdleTimeout(time.Duration(*conf.maxBackoff)))
 		dialOpts = append(dialOpts, extgrpc.EndpointGroupGRPCOpts()...)
 	}
 

pkg/receive/handler.go

@@ -15,6 +15,7 @@ import (
 	"net/http"
 	"sort"
 	"strconv"
+	"strings"
 	"sync"
 	"time"
 
@@ -396,6 +397,9 @@ func (h *Handler) getStats(r *http.Request, statsByLabelName string) ([]statusap
 // Close stops the Handler.
 func (h *Handler) Close() {
 	runutil.CloseWithLogOnErr(h.logger, h.httpSrv, "receive HTTP server")
+	if err := h.peers.closeAll(); err != nil {
+		level.Error(h.logger).Log("msg", "closing peer connections failed, we might have leaked file descriptors", "err", err)
+	}
 }
 
 // Run serves the HTTP endpoints.
@@ -790,6 +794,12 @@ func (h *Handler) distributeTimeseriesToReplicas(
 			writeDestination[endpointReplica] = writeableSeries
 		}
 	}
+	if h.receiverMode == RouterOnly && len(localWrites) > 0 {
+		panic("router only mode should not have any local writes")
+	}
+	if h.receiverMode == IngestorOnly && len(remoteWrites) > 0 {
+		panic("ingestor only mode should not have any remote writes")
+	}
 	return localWrites, remoteWrites, nil
 }
 
@@ -1226,6 +1236,7 @@ func newPeerGroup(backoff backoff.Backoff, forwardDelay prometheus.Histogram, as
 }
 
 type peersContainer interface {
+	closeAll() error
 	close(string) error
 	getConnection(context.Context, string) (WriteableStoreAsyncClient, error)
 	markPeerUnavailable(string)
@@ -1273,18 +1284,37 @@ type peerGroup struct {
 	dialer func(ctx context.Context, target string, opts ...grpc.DialOption) (conn *grpc.ClientConn, err error)
 }
 
+func (p *peerGroup) closeAll() error {
+	p.m.Lock()
+	defer p.m.Unlock()
+	var closeErrors []string // Slice to store error messages
+	for addr := range p.connections {
+		if err := p.closeUnlocked(addr); err != nil {
+			// Collect error messages
+			closeErrors = append(closeErrors, fmt.Sprintf("failed to close %s: %v", addr, err))
+		}
+	}
+	if len(closeErrors) > 0 {
+		// Return aggregated error messages
+		return fmt.Errorf("multiple errors closing connections: %s", strings.Join(closeErrors, ", "))
+	}
+	return nil
+}
+
 func (p *peerGroup) close(addr string) error {
 	p.m.Lock()
 	defer p.m.Unlock()
+	return p.closeUnlocked(addr)
+}
 
+func (p *peerGroup) closeUnlocked(addr string) error {
 	c, ok := p.connections[addr]
 	if !ok {
 		// NOTE(GiedriusS): this could be valid case when the connection
 		// was never established.
 		return nil
 	}
-
-	p.connections[addr].wp.Close()
+	c.wp.Close()
 	delete(p.connections, addr)
 	if err := c.cc.Close(); err != nil {
 		return fmt.Errorf("closing connection for %s", addr)

pkg/receive/handler_test.go

@@ -27,6 +27,7 @@ import (
 
 	"github.com/alecthomas/units"
 	"github.com/go-kit/log"
+	"github.com/go-kit/log/level"
 	"github.com/gogo/protobuf/proto"
 	"github.com/golang/snappy"
 	"github.com/pkg/errors"
@@ -40,19 +41,65 @@ import (
 	"github.com/prometheus/prometheus/storage"
 	"github.com/prometheus/prometheus/tsdb"
 	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/resolver"
 
 	"github.com/efficientgo/core/testutil"
 
 	"github.com/thanos-io/thanos/pkg/block/metadata"
+	"github.com/thanos-io/thanos/pkg/component"
 	"github.com/thanos-io/thanos/pkg/extkingpin"
 	"github.com/thanos-io/thanos/pkg/logging"
 	"github.com/thanos-io/thanos/pkg/runutil"
+	grpcserver "github.com/thanos-io/thanos/pkg/server/grpc"
+	"github.com/thanos-io/thanos/pkg/store"
 	"github.com/thanos-io/thanos/pkg/store/labelpb"
 	"github.com/thanos-io/thanos/pkg/store/storepb"
 	"github.com/thanos-io/thanos/pkg/store/storepb/prompb"
 	"github.com/thanos-io/thanos/pkg/tenancy"
 )
 
+const dnsScheme = "dns"
+
+type dnsResolver struct {
+	logger    log.Logger
+	target    resolver.Target
+	cc        resolver.ClientConn
+	addrStore map[string][]string
+}
+
+func (r *dnsResolver) start() {
+	addrStrs := r.addrStore[r.target.Endpoint()]
+	addrs := make([]resolver.Address, len(addrStrs))
+	for i, s := range addrStrs {
+		addrs[i] = resolver.Address{Addr: s}
+	}
+	if err := r.cc.UpdateState(resolver.State{Addresses: addrs}); err != nil {
+		level.Error(r.logger).Log("msg", "failed to update state", "err", err)
+	}
+}
+
+func (*dnsResolver) ResolveNow(_ resolver.ResolveNowOptions) {}
+
+func (*dnsResolver) Close() {}
+
+type dnsResolverBuilder struct {
+	logger    log.Logger
+	addrStore map[string][]string
+}
+
+func (b *dnsResolverBuilder) Build(target resolver.Target, cc resolver.ClientConn, opts resolver.BuildOptions) (resolver.Resolver, error) {
+	r := &dnsResolver{
+		logger:    b.logger,
+		target:    target,
+		cc:        cc,
+		addrStore: b.addrStore,
+	}
+	r.start()
+	return r, nil
+}
+func (*dnsResolverBuilder) Scheme() string { return dnsScheme }
+
 type fakeTenantAppendable struct {
 	f *fakeAppendable
 }
@@ -191,6 +238,10 @@ func (g *fakePeersGroup) close(addr string) error {
 	return nil
 }
 
+func (g *fakePeersGroup) closeAll() error {
+	return nil
+}
+
 func (g *fakePeersGroup) getConnection(_ context.Context, addr string) (WriteableStoreAsyncClient, error) {
 	c, ok := g.clients[addr]
 	if !ok {
@@ -1735,3 +1786,92 @@ func TestHandlerFlippingHashrings(t *testing.T) {
 	cancel()
 	wg.Wait()
 }
+
+func TestIngestorRestart(t *testing.T) {
+	var err error
+	logger := log.NewLogfmtLogger(os.Stderr)
+	addr1, addr2, addr3 := "localhost:14090", "localhost:14091", "localhost:14092"
+	ing1, ing2 := startIngestor(logger, addr1, 0), startIngestor(logger, addr2, 0)
+	defer ing1.Shutdown(err) // srv1 is stable and will only be closed after the test ends
+
+	clientAddr := "ingestor.com"
+	dnsBuilder := &dnsResolverBuilder{
+		logger:    logger,
+		addrStore: map[string][]string{clientAddr: {addr2}},
+	}
+	resolver.Register(dnsBuilder)
+	dialOpts := []grpc.DialOption{
+		grpc.WithIdleTimeout(1 * time.Second), // set idle timeout to 1s will re-establish the connection quickly
+		grpc.WithTransportCredentials(insecure.NewCredentials()),
+		grpc.WithResolvers(resolver.Get(dnsScheme)),
+	}
+	client := NewHandler(logger, &Options{
+		MaxBackoff:        1 * time.Second,
+		DialOpts:          dialOpts,
+		ReplicationFactor: 2,
+		ReceiverMode:      RouterOnly,
+		ForwardTimeout:    15 * time.Second,
+	})
+	// one of the endpoints is DNS and wire up to different backend address on the fly
+	client.Hashring(&simpleHashring{addr1, fmt.Sprintf("%s:///%s", dnsScheme, clientAddr)})
+	defer client.Close()
+
+	ctx := context.TODO()
+	data := &prompb.WriteRequest{
+		Timeseries: []prompb.TimeSeries{
+			{
+				Labels:  labelpb.ZLabelsFromPromLabels(labels.FromStrings("foo", addr3)),
+				Samples: []prompb.Sample{{Timestamp: time.Now().Unix(), Value: 123}},
+			},
+		},
+	}
+
+	err = client.handleRequest(ctx, 0, "test", data)
+	require.NoError(t, err)
+
+	// close srv2 to simulate ingestor down
+	ing2.Shutdown(err)
+	ing3 := startIngestor(logger, addr3, 2*time.Second)
+	defer ing3.Shutdown(err)
+	// bind the new backend to the same DNS
+	dnsBuilder.addrStore[clientAddr] = []string{addr3}
+
+	iter, errs := 10, 0
+	for i := 0; i < iter; i++ {
+		err = client.handleRequest(ctx, 0, "test", data)
+		if err != nil {
+			require.Error(t, errUnavailable, err)
+			errs++
+		} else {
+			break
+		}
+		time.Sleep(500 * time.Millisecond)
+	}
+	require.Greater(t, errs, 0, "expected to have unavailable errors initially")
+	require.Less(t, errs, iter, "expected to recover quickly after server restarts")
+}
+
+type fakeStoreServer struct {
+	logger log.Logger
+}
+
+func (f *fakeStoreServer) RemoteWrite(_ context.Context, in *storepb.WriteRequest) (*storepb.WriteResponse, error) {
+	level.Debug(f.logger).Log("msg", "received remote write request", "request", in.String())
+	return &storepb.WriteResponse{}, nil
+}
+
+func startIngestor(logger log.Logger, serverAddress string, delay time.Duration) *grpcserver.Server {
+	h := &fakeStoreServer{logger: logger}
+	srv := grpcserver.TestServer(logger, component.Receive, serverAddress,
+		grpcserver.WithServer(store.RegisterWritableStoreServer(h)),
+	)
+	go func() {
+		if delay > 0 {
+			time.Sleep(delay)
+		}
+		if err := srv.ListenAndServe(); err != nil {
+			level.Error(logger).Log("msg", "server error", "addr", serverAddress, "err", err)
+		}
+	}()
+	return srv
+}

pkg/server/grpc/grpc.go

@@ -163,3 +163,29 @@ func (s *Server) Shutdown(err error) {
 	}
 	level.Info(s.logger).Log("msg", "internal server is shutdown gracefully", "err", err)
 }
+
+// TestServer is used only for unit tests to create a new gRPC server.
+func TestServer(logger log.Logger, comp component.Component, address string, opts ...Option) *Server {
+	logger = log.With(logger, "service", "gRPC/server", "component", comp.String())
+	options := options{
+		network: "tcp",
+		listen:  address,
+		grpcOpts: []grpc.ServerOption{
+			grpc.MaxSendMsgSize(math.MaxInt32),
+		},
+	}
+	for _, o := range opts {
+		o.apply(&options)
+	}
+	s := grpc.NewServer(options.grpcOpts...)
+	// Register all configured servers.
+	for _, f := range options.registerServerFuncs {
+		f(s)
+	}
+	return &Server{
+		logger: logger,
+		comp:   comp,
+		srv:    s,
+		opts:   options,
+	}
+}