From 336e825e3a74161682617c578bb8a15d09668ba9 Mon Sep 17 00:00:00 2001
From: 9seconds <nineseconds@yandex.ru>
Date: Wed, 24 Mar 2021 10:16:50 +0300
Subject: [PATCH] Add wg to observer

This is to ensure a case when we still can have observer hanging around
---
 mtglib/internal/relay/relay.go | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/mtglib/internal/relay/relay.go b/mtglib/internal/relay/relay.go
index 2785b03..2a27708 100644
--- a/mtglib/internal/relay/relay.go
+++ b/mtglib/internal/relay/relay.go
@@ -34,10 +34,10 @@ func (r *Relay) Process(eastConn, westConn io.ReadWriteCloser) error {
 		westConn.Close()
 	}()
 
-	go r.runObserver(r.ctx)
-
 	wg := &sync.WaitGroup{}
-	wg.Add(2) // nolint: gomnd
+	wg.Add(3) // nolint: gomnd
+
+	go r.runObserver(r.ctx, wg)
 
 	go r.transmit(eastConn, westConn, r.westBuffer, "west", wg)
 
@@ -72,7 +72,7 @@ func (r *Relay) transmit(src io.ReadCloser, dst io.WriteCloser,
 	}
 }
 
-func (r *Relay) runObserver(ctx context.Context) {
+func (r *Relay) runObserver(ctx context.Context, wg *sync.WaitGroup) {
 	ticker := time.NewTicker(time.Second)
 
 	defer func() {
@@ -82,6 +82,8 @@ func (r *Relay) runObserver(ctx context.Context) {
 		case <-ticker.C:
 		default:
 		}
+
+		wg.Done()
 	}()
 
 	lastTickAt := time.Now()