From b524da03f23260040492f68f620f2c55af4dc74c Mon Sep 17 00:00:00 2001
From: Andrew Kroh <andrew.kroh@elastic.co>
Date: Mon, 29 Jun 2020 15:45:43 -0400
Subject: [PATCH] Use ctxtool.WithFunc

---
 CHANGELOG.next.asciidoc                 |  2 +-
 filebeat/inputsource/common/listener.go | 11 ++++-------
 2 files changed, 5 insertions(+), 8 deletions(-)

diff --git a/CHANGELOG.next.asciidoc b/CHANGELOG.next.asciidoc
index aeffe760f77..51ad3eaea8b 100644
--- a/CHANGELOG.next.asciidoc
+++ b/CHANGELOG.next.asciidoc
@@ -196,7 +196,7 @@ https://github.com/elastic/beats/compare/v7.0.0-alpha2...master[Check the HEAD d
 - Fix `netflow` module to support 7 bytepad for IPFIX template. {issue}18098[18098]
 - Fix Cisco ASA dissect pattern for 313008 & 313009 messages. {pull}19149[19149]
 - Fix date and timestamp formats for fortigate module {pull}19316[19316]
-- Fix memory leak in tcp and unix imput sources. {pull}19459[19459]
+- Fix memory leak in tcp and unix input sources. {pull}19459[19459]
 
 *Heartbeat*
 
diff --git a/filebeat/inputsource/common/listener.go b/filebeat/inputsource/common/listener.go
index 4692e474930..c3ed0362887 100644
--- a/filebeat/inputsource/common/listener.go
+++ b/filebeat/inputsource/common/listener.go
@@ -27,6 +27,7 @@ import (
 
 	"github.com/elastic/beats/v7/libbeat/common/atomic"
 	"github.com/elastic/beats/v7/libbeat/logp"
+	"github.com/elastic/go-concert/ctxtool"
 )
 
 // Family represents the type of connection we're handling
@@ -113,17 +114,12 @@ func (l *Listener) run() {
 			}
 		}
 
-		handler := l.handlerFactory(*l.config)
-		ctx, cancel := context.WithCancel(l.ctx)
-		go func() {
-			<-ctx.Done()
-			conn.Close()
-		}()
-
 		l.wg.Add(1)
 		go func() {
 			defer logp.Recover("recovering from a " + l.family.String() + " client crash")
 			defer l.wg.Done()
+
+			ctx, cancel := ctxtool.WithFunc(l.ctx, func() { conn.Close() })
 			defer cancel()
 
 			l.registerHandler()
@@ -136,6 +132,7 @@ func (l *Listener) run() {
 				l.log.Debugw("New client", "remote_address", conn.RemoteAddr(), "total", l.clientsCount.Load())
 			}
 
+			handler := l.handlerFactory(*l.config)
 			err := handler(ctx, conn)
 			if err != nil {
 				l.log.Debugw("client error", "error", err)