From ba162700592877bfa098bdaafc7349f0c70d0eaa Mon Sep 17 00:00:00 2001
From: "Quentin McGaw (desktop)" <quentin.mcgaw@gmail.com>
Date: Mon, 16 Aug 2021 19:19:33 +0000
Subject: [PATCH] Maint: context aware collectLines functions

---
 internal/dns/logs.go     | 19 +++++++++++++------
 internal/dns/setup.go    | 11 ++++++-----
 internal/openvpn/logs.go | 21 +++++++++++++--------
 internal/openvpn/run.go  |  7 ++++---
 4 files changed, 36 insertions(+), 22 deletions(-)

diff --git a/internal/dns/logs.go b/internal/dns/logs.go
index 53347f7b..1acb300c 100644
--- a/internal/dns/logs.go
+++ b/internal/dns/logs.go
@@ -1,6 +1,7 @@
 package dns
 
 import (
+	"context"
 	"regexp"
 	"strings"
 
@@ -8,18 +9,24 @@ import (
 	"github.com/qdm12/golibs/logging"
 )
 
-func (l *Loop) collectLines(stdout, stderr <-chan string, done chan<- struct{}) {
+func (l *Loop) collectLines(ctx context.Context, done chan<- struct{},
+	stdout, stderr chan string) {
 	defer close(done)
+
 	var line string
-	var ok bool
+
 	for {
 		select {
-		case line, ok = <-stderr:
-		case line, ok = <-stdout:
-		}
-		if !ok {
+		case <-ctx.Done():
+			// Context should only be canceled after stdout and stderr are done
+			// being written to.
+			close(stdout)
+			close(stderr)
 			return
+		case line = <-stderr:
+		case line = <-stdout:
 		}
+
 		line, level := processLogLine(line)
 		switch level {
 		case logging.LevelDebug:
diff --git a/internal/dns/setup.go b/internal/dns/setup.go
index 1bb84f74..11079faa 100644
--- a/internal/dns/setup.go
+++ b/internal/dns/setup.go
@@ -30,12 +30,13 @@ func (l *Loop) setupUnbound(ctx context.Context) (
 		return nil, nil, nil, err
 	}
 
-	collectLinesDone := make(chan struct{})
-	go l.collectLines(stdoutLines, stderrLines, collectLinesDone)
+	linesCollectionCtx, linesCollectionCancel := context.WithCancel(context.Background())
+	lineCollectionDone := make(chan struct{})
+	go l.collectLines(linesCollectionCtx, lineCollectionDone,
+		stdoutLines, stderrLines)
 	closeStreams = func() {
-		close(stdoutLines)
-		close(stderrLines)
-		<-collectLinesDone
+		linesCollectionCancel()
+		<-lineCollectionDone
 	}
 
 	// use Unbound
diff --git a/internal/openvpn/logs.go b/internal/openvpn/logs.go
index f7ab7181..7db76477 100644
--- a/internal/openvpn/logs.go
+++ b/internal/openvpn/logs.go
@@ -1,6 +1,7 @@
 package openvpn
 
 import (
+	"context"
 	"strings"
 
 	"github.com/fatih/color"
@@ -8,20 +9,24 @@ import (
 	"github.com/qdm12/golibs/logging"
 )
 
-func (l *Loop) collectLines(stdout, stderr <-chan string, done chan<- struct{}) {
+func (l *Loop) collectLines(ctx context.Context, done chan<- struct{},
+	stdout, stderr chan string) {
 	defer close(done)
+
 	var line string
-	var ok, errLine bool
 
 	for {
-		errLine = false
+		errLine := false
 		select {
-		case line, ok = <-stdout:
-		case line, ok = <-stderr:
-			errLine = true
-		}
-		if !ok {
+		case <-ctx.Done():
+			// Context should only be canceled after stdout and stderr are done
+			// being written to.
+			close(stdout)
+			close(stderr)
 			return
+		case line = <-stdout:
+		case line = <-stderr:
+			errLine = true
 		}
 		line, level := processLogLine(line)
 		if line == "" {
diff --git a/internal/openvpn/run.go b/internal/openvpn/run.go
index 5e44e9a2..f9a607c3 100644
--- a/internal/openvpn/run.go
+++ b/internal/openvpn/run.go
@@ -80,11 +80,12 @@ func (l *Loop) Run(ctx context.Context, done chan<- struct{}) {
 			continue
 		}
 
+		linesCollectionCtx, linesCollectionCancel := context.WithCancel(context.Background())
 		lineCollectionDone := make(chan struct{})
-		go l.collectLines(stdoutLines, stderrLines, lineCollectionDone)
+		go l.collectLines(linesCollectionCtx, lineCollectionDone,
+			stdoutLines, stderrLines)
 		closeStreams := func() {
-			close(stdoutLines)
-			close(stderrLines)
+			linesCollectionCancel()
 			<-lineCollectionDone
 		}