From 198d4974e312969924f7ae2e23b62cd40f6693af Mon Sep 17 00:00:00 2001
From: Bradley Cicenas <bradley.cicenas@gmail.com>
Date: Sun, 5 Feb 2017 00:25:30 +0000
Subject: [PATCH] add waitgroup, exit handlers to logging server

---
 logging/main.go | 34 ++++++++++++++++++----------------
 main.go         |  2 ++
 2 files changed, 20 insertions(+), 16 deletions(-)

diff --git a/logging/main.go b/logging/main.go
index 4cf970f..adf79e4 100644
--- a/logging/main.go
+++ b/logging/main.go
@@ -3,6 +3,7 @@ package logging
 import (
 	"fmt"
 	"net"
+	"sync"
 	"time"
 
 	"github.com/op/go-logging"
@@ -16,23 +17,23 @@ const (
 var format = logging.MustStringFormatter(
 	`%{color}%{time:15:04:05.000} %{shortfunc} ▶ %{level:.4s} %{id:03x}%{color:reset} %{message}`,
 )
+var exited bool
+var wg sync.WaitGroup
 
 type CTopLogger struct {
 	*logging.Logger
 	backend *logging.MemoryBackend
-	done    chan bool
 }
 
 func New(serverEnabled string) *CTopLogger {
 	log := &CTopLogger{
 		logging.MustGetLogger("ctop"),
 		logging.NewMemoryBackend(size),
-		make(chan bool),
 	}
 
 	backendFmt := logging.NewBackendFormatter(log.backend, format)
 	logging.SetBackend(backendFmt)
-	log.Info("logger initialized")
+	log.Notice("logger initialized")
 
 	if serverEnabled == "1" {
 		log.Serve()
@@ -42,7 +43,8 @@ func New(serverEnabled string) *CTopLogger {
 }
 
 func (log *CTopLogger) Exit() {
-	log.done <- true
+	exited = true
+	wg.Wait()
 }
 
 func (log *CTopLogger) Serve() {
@@ -51,35 +53,31 @@ func (log *CTopLogger) Serve() {
 		panic(err)
 	}
 
-	go func() {
-		switch {
-		case <-log.done:
-			ln.Close()
-			return
-		default:
-			//
-		}
-	}()
-
 	go func() {
 		for {
 			conn, err := ln.Accept()
 			if err != nil {
-				panic(err)
+				if err, ok := err.(net.Error); ok && err.Temporary() {
+					continue
+				}
+				return
 			}
 			go log.handler(conn)
 		}
 	}()
 
-	log.Info("logging server started")
+	log.Notice("logging server started")
 }
 
 func (log *CTopLogger) handler(conn net.Conn) {
+	wg.Add(1)
+	defer wg.Done()
 	defer conn.Close()
 	for msg := range log.tail() {
 		msg = fmt.Sprintf("%s\n", msg)
 		conn.Write([]byte(msg))
 	}
+	conn.Write([]byte("bye\n"))
 }
 
 func (log *CTopLogger) tail() chan string {
@@ -95,6 +93,10 @@ func (log *CTopLogger) tail() chan string {
 					node = nnode
 					break
 				}
+				if exited {
+					close(stream)
+					return
+				}
 				time.Sleep(1 * time.Second)
 			}
 		}
diff --git a/main.go b/main.go
index 111b95c..d46df6b 100644
--- a/main.go
+++ b/main.go
@@ -17,6 +17,8 @@ func main() {
 	for {
 		exit := Display(g)
 		if exit {
+			log.Notice("shutting down")
+			log.Exit()
 			return
 		}
 	}