From ce4b9976236d1fc6091fa5270a3f6518a08a8762 Mon Sep 17 00:00:00 2001
From: Yaroslav Vorobiov <yar.vorobiov@gmail.com>
Date: Thu, 4 Jul 2019 13:34:23 +0300
Subject: [PATCH] storagenode/nodestats: connection leak (#2443)

---
 storagenode/nodestats/service.go | 37 +++++++++++++++++++++++++++-----
 storagenode/orders/sender.go     |  2 +-
 2 files changed, 33 insertions(+), 6 deletions(-)

diff --git a/storagenode/nodestats/service.go b/storagenode/nodestats/service.go
index c6cb5dd19..f3964a2cf 100644
--- a/storagenode/nodestats/service.go
+++ b/storagenode/nodestats/service.go
@@ -9,6 +9,7 @@ import (
 
 	"github.com/zeebo/errs"
 	"go.uber.org/zap"
+	"google.golang.org/grpc"
 	"gopkg.in/spacemonkeygo/monkit.v2"
 
 	"storj.io/storj/pkg/kademlia"
@@ -42,6 +43,17 @@ type SpaceUsageStamp struct {
 	TimeStamp time.Time
 }
 
+// Client encapsulates NodeStatsClient with underlying connection
+type Client struct {
+	conn *grpc.ClientConn
+	pb.NodeStatsClient
+}
+
+// Close closes underlying client connection
+func (c *Client) Close() error {
+	return c.conn.Close()
+}
+
 // Service retrieves info from satellites using GRPC client
 type Service struct {
 	log *zap.Logger
@@ -63,11 +75,17 @@ func NewService(log *zap.Logger, transport transport.Client, kademlia *kademlia.
 func (s *Service) GetUptimeCheckForSatellite(ctx context.Context, satelliteID storj.NodeID) (_ *UptimeCheck, err error) {
 	defer mon.Task()(&ctx)(&err)
 
-	client, err := s.getGRPCClientForSatellite(ctx, satelliteID)
+	client, err := s.DialNodeStats(ctx, satelliteID)
 	if err != nil {
 		return nil, NodeStatsServiceErr.Wrap(err)
 	}
 
+	defer func() {
+		if cerr := client.Close(); cerr != nil {
+			err = errs.Combine(err, NodeStatsServiceErr.New("failed to close connection: %v", cerr))
+		}
+	}()
+
 	resp, err := client.UptimeCheck(ctx, &pb.UptimeCheckRequest{})
 	if err != nil {
 		return nil, NodeStatsServiceErr.Wrap(err)
@@ -86,11 +104,17 @@ func (s *Service) GetUptimeCheckForSatellite(ctx context.Context, satelliteID st
 func (s *Service) GetDailyStorageUsedForSatellite(ctx context.Context, satelliteID storj.NodeID, from, to time.Time) (_ []SpaceUsageStamp, err error) {
 	defer mon.Task()(&ctx)(&err)
 
-	client, err := s.getGRPCClientForSatellite(ctx, satelliteID)
+	client, err := s.DialNodeStats(ctx, satelliteID)
 	if err != nil {
 		return nil, NodeStatsServiceErr.Wrap(err)
 	}
 
+	defer func() {
+		if cerr := client.Close(); cerr != nil {
+			err = errs.Combine(err, NodeStatsServiceErr.New("failed to close connection: %v", cerr))
+		}
+	}()
+
 	resp, err := client.DailyStorageUsage(ctx, &pb.DailyStorageUsageRequest{})
 	if err != nil {
 		return nil, NodeStatsServiceErr.Wrap(err)
@@ -99,8 +123,8 @@ func (s *Service) GetDailyStorageUsedForSatellite(ctx context.Context, satellite
 	return fromSpaceUsageResponse(resp, satelliteID), nil
 }
 
-// getGRPCClientForSatellite inits GRPC client for the satellite by id
-func (s *Service) getGRPCClientForSatellite(ctx context.Context, satelliteID storj.NodeID) (pb.NodeStatsClient, error) {
+// DialNodeStats dials GRPC NodeStats client for the satellite by id
+func (s *Service) DialNodeStats(ctx context.Context, satelliteID storj.NodeID) (*Client, error) {
 	satellite, err := s.kademlia.FindNode(ctx, satelliteID)
 	if err != nil {
 		return nil, errs.New("unable to find satellite %s: %v", satelliteID, err)
@@ -111,7 +135,10 @@ func (s *Service) getGRPCClientForSatellite(ctx context.Context, satelliteID sto
 		return nil, errs.New("unable to connect to the satellite %s: %v", satelliteID, err)
 	}
 
-	return pb.NewNodeStatsClient(conn), nil
+	return &Client{
+		conn:            conn,
+		NodeStatsClient: pb.NewNodeStatsClient(conn),
+	}, nil
 }
 
 // fromSpaceUsageResponse get SpaceUsageStamp slice from pb.SpaceUsageResponse
diff --git a/storagenode/orders/sender.go b/storagenode/orders/sender.go
index e4f5837ce..76e29393c 100644
--- a/storagenode/orders/sender.go
+++ b/storagenode/orders/sender.go
@@ -165,7 +165,7 @@ func (sender *Sender) settle(ctx context.Context, log *zap.Logger, satelliteID s
 	}
 	defer func() {
 		if cerr := conn.Close(); cerr != nil {
-			err = errs.Combine(err, OrderError.New("failed to close connection: %v", err))
+			err = errs.Combine(err, OrderError.New("failed to close connection: %v", cerr))
 		}
 	}()