[ES-1122621] Fix receive connection idle if ingestor restarts

Signed-off-by: Yi Jin <[email protected]>

Author: jnyi

pkg/receive/handler.go

@@ -149,7 +149,8 @@ func NewHandler(logger log.Logger, o *Options) *Handler {
 		writer:  o.Writer,
 		router:  route.New(),
 		options: o,
-		peers: newPeerGroup(logger,
+		peers: newPeerGroup(
+			logger,
 			backoff.Backoff{
 				Factor: 2,
 				Min:    100 * time.Millisecond,
@@ -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.
@@ -1218,6 +1222,27 @@ type peerWorker struct {
 	forwardDelay prometheus.Histogram
 }
 
+func (pw *peerWorker) isReady() bool {
+	return pw.cc != nil
+}
+
+func (pw *peerWorker) markNotReady() error {
+	if pw.cc != nil {
+		err := pw.cc.Close()
+		pw.cc = nil
+		return err
+	}
+	return nil
+}
+
+func (pw *peerWorker) close() error {
+	pw.wp.Close()
+	if pw.cc != nil {
+		return pw.cc.Close()
+	}
+	return nil
+}
+
 func newPeerGroup(logger log.Logger, backoff backoff.Backoff, forwardDelay prometheus.Histogram, asyncForwardWorkersCount uint, dialOpts ...grpc.DialOption) peersContainer {
 	return &peerGroup{
 		logger:                   logger,
@@ -1233,6 +1258,7 @@ func newPeerGroup(logger log.Logger, backoff backoff.Backoff, forwardDelay prome
 }
 
 type peersContainer interface {
+	closeAll() error
 	close(string) error
 	getConnection(context.Context, string) (WriteableStoreAsyncClient, error)
 	markPeerUnavailable(string)
@@ -1281,6 +1307,17 @@ 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()
+	for addr := range p.connections {
+		if err := p.closeUnlocked(addr); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
 func (p *peerGroup) close(addr string) error {
 	p.m.Lock()
 	defer p.m.Unlock()
@@ -1294,10 +1331,8 @@ func (p *peerGroup) closeUnlocked(addr string) error {
 		// was never established.
 		return nil
 	}
-	level.Error(p.logger).Log("msg", "closing connection", "addr", addr)
-	p.connections[addr].wp.Close()
 	delete(p.connections, addr)
-	if err := c.cc.Close(); err != nil {
+	if err := c.close(); err != nil {
 		return fmt.Errorf("closing connection for %s", addr)
 	}
 
@@ -1312,6 +1347,7 @@ func (p *peerGroup) getConnection(ctx context.Context, addr string) (WriteableSt
 	// use a RLock first to prevent blocking if we don't need to.
 	p.m.RLock()
 	c, ok := p.connections[addr]
+	ok = ok && c.isReady()
 	p.m.RUnlock()
 	if ok {
 		return c, nil
@@ -1321,36 +1357,32 @@ func (p *peerGroup) getConnection(ctx context.Context, addr string) (WriteableSt
 	defer p.m.Unlock()
 	// Make sure that another caller hasn't created the connection since obtaining the write lock.
 	c, ok = p.connections[addr]
-	if ok {
+	if ok && c.isReady() {
 		return c, nil
 	}
-	level.Info(p.logger).Log("msg", "dialing peer", "peer", addr)
+	level.Debug(p.logger).Log("msg", "dialing peer", "addr", addr)
 	conn, err := p.dialer(ctx, addr, p.dialOpts...)
 	if err != nil {
-		// clear retry state if dial failed, this is necessary to avoid dialing peer too often.
-		delete(p.peerStates, addr)
 		p.markPeerUnavailableUnlocked(addr)
 		dialError := errors.Wrap(err, "failed to dial peer")
 		return nil, errors.Wrap(dialError, errUnavailable.Error())
 	}
-
-	p.connections[addr] = newPeerWorker(conn, p.forwardDelay, p.asyncForwardWorkersCount)
+	if !ok {
+		p.connections[addr] = newPeerWorker(conn, p.forwardDelay, p.asyncForwardWorkersCount)
+	} else {
+		p.connections[addr].cc = conn
+	}
 	return p.connections[addr], nil
 }
 
 func (p *peerGroup) markPeerUnavailable(addr string) {
 	p.m.Lock()
 	defer p.m.Unlock()
 
-	if p.markPeerUnavailableUnlocked(addr) {
-		// close the connection if the backoff tries too many times, likely server is down.
-		if err := p.closeUnlocked(addr); err != nil {
-			level.Error(p.logger).Log("msg", "failed to close connection", "addr", addr, "err", err)
-		}
-	}
+	p.markPeerUnavailableUnlocked(addr)
 }
 
-func (p *peerGroup) markPeerUnavailableUnlocked(addr string) bool {
+func (p *peerGroup) markPeerUnavailableUnlocked(addr string) {
 	state, ok := p.peerStates[addr]
 	if !ok {
 		state = &retryState{attempt: -1}
@@ -1359,7 +1391,12 @@ func (p *peerGroup) markPeerUnavailableUnlocked(addr string) bool {
 	delay := p.expBackoff.ForAttempt(state.attempt)
 	state.nextAllowed = time.Now().Add(delay)
 	p.peerStates[addr] = state
-	return delay >= p.expBackoff.Max
+	if delay >= p.expBackoff.Max {
+		level.Warn(p.logger).Log("msg", "peer is not ready", "addr", addr, "next_allowed", state.nextAllowed)
+		if err := p.connections[addr].markNotReady(); err != nil {
+			level.Error(p.logger).Log("msg", "failed to mark peer as not ready", "err", err)
+		}
+	}
 }
 
 func (p *peerGroup) markPeerAvailable(addr string) {

pkg/receive/handler_test.go

@@ -26,7 +26,9 @@ import (
 	"gopkg.in/yaml.v3"
 
 	"github.com/alecthomas/units"
+	"github.com/efficientgo/core/testutil"
 	"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 +42,63 @@ import (
 	"github.com/prometheus/prometheus/storage"
 	"github.com/prometheus/prometheus/tsdb"
 	"github.com/stretchr/testify/require"
-
-	"github.com/efficientgo/core/testutil"
+	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/resolver"
 
 	"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 +237,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 {
@@ -1736,50 +1786,90 @@ func TestHandlerFlippingHashrings(t *testing.T) {
 	wg.Wait()
 }
 
-func TestPeerGroup(t *testing.T) {
+func TestIngestorRestart(t *testing.T) {
+	var err error
 	logger := log.NewLogfmtLogger(os.Stderr)
-	serverAddress := "http://localhost:19090"
-	dialOpts := []grpc.DialOption{grpc.WithInsecure()}
-	srv := startServer(logger, serverAddress)
-	ctx := context.TODO()
+	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.WithTransportCredentials(insecure.NewCredentials()),
+		grpc.WithResolvers(resolver.Get(dnsScheme)),
+	}
 	client := NewHandler(logger, &Options{
-		MaxBackoff: 1 * time.Second,
-		DialOpts:   dialOpts,
+		MaxBackoff:        1 * time.Second,
+		DialOpts:          dialOpts,
+		ReplicationFactor: 2,
+		ReceiverMode:      RouterOnly,
+		ForwardTimeout:    15 * time.Second,
 	})
-	_, err := client.peers.getConnection(ctx, serverAddress)
-	require.NoError(t, err)
-	// close the server and wait for the backoff to kick in and see how long it takes
-	srv.Close()
-	// server is closed, now we can't send requests to it
-	er := endpointReplica{endpoint: serverAddress, replica: 0}
-	data := trackedSeries{
-		timeSeries: []prompb.TimeSeries{
+	// 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", "bar")),
+				Labels:  labelpb.ZLabelsFromPromLabels(labels.FromStrings("foo", addr3)),
 				Samples: []prompb.Sample{{Timestamp: time.Now().Unix(), Value: 123}},
 			},
 		},
 	}
 
-	N := 50
-	responses := make(chan writeResponse, N)
-	for i := 0; i < N; i++ {
-		var wg sync.WaitGroup
-		wg.Add(1)
-		client.sendRemoteWrite(ctx, "", er, data, false, responses, &wg)
-		wg.Wait()
-		_, err = client.peers.getConnection(ctx, serverAddress)
-		require.Error(t, errUnavailable, err)
-		time.Sleep(100 * time.Millisecond)
+	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")
 }
 
-func startServer(logger log.Logger, serverAddress string) *Handler {
-	srv := NewHandler(logger, &Options{
-		ListenAddress: serverAddress,
-	})
+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() {
-		srv.Run()
+		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,
+	}
+}