Maint: context aware collectLines functions

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 <-ctx.Done():
-		case line, ok = <-stdout:
+			// Context should only be canceled after stdout and stderr are done
-		}
+			// being written to.
-		if !ok {
+			close(stdout)
+			close(stderr)
 			return
+		case line = <-stderr:
+		case line = <-stdout:
 		}
+
 		line, level := processLogLine(line)
 		switch level {
 		case logging.LevelDebug:

internal/dns/setup.go

@@ -30,12 +30,13 @@ func (l *Loop) setupUnbound(ctx context.Context) (
 		return nil, nil, nil, err
 	}
 
-	collectLinesDone := make(chan struct{})
+	linesCollectionCtx, linesCollectionCancel := context.WithCancel(context.Background())
-	go l.collectLines(stdoutLines, stderrLines, collectLinesDone)
+	lineCollectionDone := make(chan struct{})
+	go l.collectLines(linesCollectionCtx, lineCollectionDone,
+		stdoutLines, stderrLines)
 	closeStreams = func() {
-		close(stdoutLines)
+		linesCollectionCancel()
-		close(stderrLines)
+		<-lineCollectionDone
-		<-collectLinesDone
 	}
 
 	// use Unbound

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 <-ctx.Done():
-		case line, ok = <-stderr:
+			// Context should only be canceled after stdout and stderr are done
-			errLine = true
+			// being written to.
-		}
+			close(stdout)
-		if !ok {
+			close(stderr)
 			return
+		case line = <-stdout:
+		case line = <-stderr:
+			errLine = true
 		}
 		line, level := processLogLine(line)
 		if line == "" {

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)
+			linesCollectionCancel()
-			close(stderrLines)
 			<-lineCollectionDone
 		}