[shell-operator] Close resources correctly with dcql -w (#798)

Signed-off-by: kerimov <brileyforbusiness@gmail.com>

Author: brileyyyy

pkg/debug/client.go

@@ -7,54 +7,58 @@ import (
 	"io"
 	"net"
 	"net/http"
+	"time"
 
 	"github.com/flant/shell-operator/pkg/app"
 	utils "github.com/flant/shell-operator/pkg/utils/file"
 )
 
 type Client struct {
 	SocketPath string
+	httpClient *http.Client
 }
 
-func NewClient() *Client {
-	return &Client{}
-}
-
-func (c *Client) WithSocketPath(path string) {
-	c.SocketPath = path
-}
-
-func (c *Client) newHttpClient() (http.Client, error) {
-	exists, err := utils.FileExists(c.SocketPath)
+func NewClient(socketPath string) (*Client, error) {
+	exists, err := utils.FileExists(socketPath)
 	if err != nil {
-		return http.Client{}, fmt.Errorf("check debug socket '%s': %s", c.SocketPath, err)
+		return nil, fmt.Errorf("check debug socket '%s': %s", socketPath, err)
 	}
 	if !exists {
-		return http.Client{}, fmt.Errorf("debug socket '%s' is not exists", c.SocketPath)
+		return nil, fmt.Errorf("debug socket '%s' is not exists", socketPath)
 	}
 
-	return http.Client{
+	client := &http.Client{
 		Transport: &http.Transport{
-			DialContext: func(_ context.Context, _, _ string) (net.Conn, error) {
-				return net.Dial("unix", c.SocketPath)
+			DialContext: func(ctx context.Context, _, _ string) (net.Conn, error) {
+				dialer := &net.Dialer{
+					Timeout: 10 * time.Second,
+				}
+				return dialer.DialContext(ctx, "unix", socketPath)
 			},
+			DisableKeepAlives: true,
 		},
+	}
+
+	return &Client{
+		SocketPath: socketPath,
+		httpClient: client,
 	}, nil
 }
 
-func DefaultClient() *Client {
-	cl := NewClient()
-	cl.WithSocketPath(app.DebugUnixSocket)
-	return cl
+func (c *Client) Close() {
+	if c.httpClient != nil {
+		if transport, ok := c.httpClient.Transport.(*http.Transport); ok {
+			transport.CloseIdleConnections()
+		}
+	}
 }
 
-func (c *Client) Get(url string) ([]byte, error) {
-	httpc, err := c.newHttpClient()
-	if err != nil {
-		return nil, err
-	}
+func DefaultClient() (*Client, error) {
+	return NewClient(app.DebugUnixSocket)
+}
 
-	resp, err := httpc.Get(url)
+func (c *Client) Get(url string) ([]byte, error) {
+	resp, err := c.httpClient.Get(url)
 	if err != nil {
 		return nil, err
 	}
@@ -69,12 +73,7 @@ func (c *Client) Get(url string) ([]byte, error) {
 }
 
 func (c *Client) Post(targetUrl string, data map[string][]string) ([]byte, error) {
-	httpc, err := c.newHttpClient()
-	if err != nil {
-		return nil, err
-	}
-
-	resp, err := httpc.PostForm(targetUrl, data)
+	resp, err := c.httpClient.PostForm(targetUrl, data)
 	if err != nil {
 		return nil, err
 	}

pkg/debug/debug-cmd.go

@@ -28,12 +28,19 @@ func DefineDebugCommands(kpApp *kingpin.Application) {
 			var refreshInterval time.Duration
 			output := termenv.NewOutput(os.Stdout)
 
+			client, err := DefaultClient()
+			if err != nil {
+				return err
+			}
+			defer client.Close()
+
 			if watch {
 				c := make(chan os.Signal, 1)
 				signal.Notify(c, os.Interrupt)
 
 				go func() {
 					<-c
+					client.Close()
 					output.ExitAltScreen()
 					os.Exit(0)
 				}()
@@ -49,10 +56,10 @@ func DefineDebugCommands(kpApp *kingpin.Application) {
 				}
 			}
 
-			client := Queue(DefaultClient())
+			queueClient := Queue(client)
 
 			for {
-				out, err := client.List(outputFormat, showEmpty)
+				out, err := queueClient.List(outputFormat, showEmpty)
 				if err != nil {
 					return err
 				}
@@ -85,7 +92,11 @@ func DefineDebugCommands(kpApp *kingpin.Application) {
 
 	queueMainCmd := queueCmd.Command("main", "Dump tasks in the main queue.").
 		Action(func(_ *kingpin.ParseContext) error {
-			out, err := Queue(DefaultClient()).Main(outputFormat)
+			client, err := DefaultClient()
+			if err != nil {
+				return err
+			}
+			out, err := Queue(client).Main(outputFormat)
 			if err != nil {
 				return err
 			}
@@ -100,7 +111,11 @@ func DefineDebugCommands(kpApp *kingpin.Application) {
 
 	configListCmd := configCmd.Command("list", "List available runtime parameters.").
 		Action(func(_ *kingpin.ParseContext) error {
-			out, err := Config(DefaultClient()).List(outputFormat)
+			client, err := DefaultClient()
+			if err != nil {
+				return err
+			}
+			out, err := Config(client).List(outputFormat)
 			if err != nil {
 				return err
 			}
@@ -115,7 +130,11 @@ func DefineDebugCommands(kpApp *kingpin.Application) {
 	var paramDuration time.Duration
 	configSetCmd := configCmd.Command("set", "Set runtime parameter.").
 		Action(func(_ *kingpin.ParseContext) error {
-			out, err := Config(DefaultClient()).Set(paramName, paramValue, paramDuration)
+			client, err := DefaultClient()
+			if err != nil {
+				return err
+			}
+			out, err := Config(client).Set(paramName, paramValue, paramDuration)
 			if err != nil {
 				return err
 			}
@@ -131,8 +150,12 @@ func DefineDebugCommands(kpApp *kingpin.Application) {
 	var rawUrl string
 	rawCommand := app.CommandWithDefaultUsageTemplate(kpApp, "raw", "Make a raw request to debug endpoint.").
 		Action(func(_ *kingpin.ParseContext) error {
+			client, err := DefaultClient()
+			if err != nil {
+				return err
+			}
 			url := fmt.Sprintf("http://unix%s", rawUrl)
-			resp, err := DefaultClient().Get(url)
+			resp, err := client.Get(url)
 			if err != nil {
 				return err
 			}
@@ -148,7 +171,11 @@ func DefineDebugCommandsSelf(kpApp *kingpin.Application) {
 	hookCmd := app.CommandWithDefaultUsageTemplate(kpApp, "hook", "Actions for hooks")
 	hookListCmd := hookCmd.Command("list", "List all hooks.").
 		Action(func(_ *kingpin.ParseContext) error {
-			outBytes, err := Hook(DefaultClient()).List(outputFormat)
+			client, err := DefaultClient()
+			if err != nil {
+				return err
+			}
+			outBytes, err := Hook(client).List(outputFormat)
 			if err != nil {
 				return err
 			}
@@ -162,7 +189,11 @@ func DefineDebugCommandsSelf(kpApp *kingpin.Application) {
 	var hookName string
 	hookSnapshotCmd := hookCmd.Command("snapshot", "Dump hook snapshots.").
 		Action(func(_ *kingpin.ParseContext) error {
-			outBytes, err := Hook(DefaultClient()).Name(hookName).Snapshots(outputFormat)
+			client, err := DefaultClient()
+			if err != nil {
+				return err
+			}
+			outBytes, err := Hook(client).Name(hookName).Snapshots(outputFormat)
 			if err != nil {
 				return err
 			}