Add logic to flush channels, and to safeguard adding of events, after main loop in Application.Run exited

application.go

@@ -1,6 +1,7 @@
 package tview
 
 import (
+	"context"
 	"sync"
 	"time"
 
@@ -66,6 +67,9 @@ type queuedUpdate struct {
 type Application struct {
 	sync.RWMutex
 
+	runContext    context.Context
+	runCancelFunc context.CancelFunc
+
 	// The application's screen. Apart from Run(), this variable should never be
 	// set directly. Always use the screenReplacement channel after calling
 	// Fini(), to set a new screen (or nil to stop the application).
@@ -120,9 +124,39 @@ type Application struct {
 	lastMouseButtons        tcell.ButtonMask // The last mouse button state.
 }
 
+func (a *Application) Close() error {
+	a.runCancelFunc()
+	close(a.events)
+	close(a.screenReplacement)
+	close(a.updates)
+
+	// flush events channel
+	go func() {
+		for range a.events {
+		}
+	}()
+	// flush screenReplacement channel
+	go func() {
+		for range a.screenReplacement {
+		}
+	}()
+	// flush updates channel
+	go func() {
+		for up := range a.updates {
+			// important  to set done for calling channel to be able to return
+			up.done <- struct{}{}
+		}
+	}()
+
+	return nil
+}
+
 // NewApplication creates and returns a new application.
 func NewApplication() *Application {
+	cancelContext, cancelFunc := context.WithCancel(context.Background())
 	return &Application{
+		runContext:        cancelContext,
+		runCancelFunc:     cancelFunc,
 		events:            make(chan tcell.Event, queueSize),
 		updates:           make(chan queuedUpdate, queueSize),
 		screenReplacement: make(chan tcell.Screen, 1),
@@ -188,7 +222,10 @@ func (a *Application) SetScreen(screen tcell.Screen) *Application {
 	oldScreen := a.screen
 	a.Unlock()
 	oldScreen.Fini()
-	a.screenReplacement <- screen
+	// check to see if the Application.Run is still valid
+	if a.runContext.Err() == nil {
+		a.screenReplacement <- screen
+	}
 
 	return a
 }
@@ -252,8 +289,15 @@ func (a *Application) Run() error {
 	var wg sync.WaitGroup
 	wg.Add(1)
 	go func() {
+		defer func() {
+			// call the runCancelFunc when exiting this function.
+			//This will stop the channels accepting any more events
+			a.runCancelFunc()
+		}()
 		defer wg.Done()
-		for {
+
+		// check to see if the Application.Run is still valid
+		for a.runContext.Err() == nil {
 			a.RLock()
 			screen := a.screen
 			a.RUnlock()
@@ -271,37 +315,47 @@ func (a *Application) Run() error {
 				continue
 			}
 
-			// A screen was finalized (event is nil). Wait for a new scren.
-			screen = <-a.screenReplacement
-			if screen == nil {
-				// No new screen. We're done.
-				a.QueueEvent(nil)
+			// A screen was finalized (event is nil). Wait for a new screen.
+			var ok bool
+			select {
+			// exit when runContext complete
+			case <-a.runContext.Done():
 				return
-			}
+			case screen, ok = <-a.screenReplacement:
+				if !ok || screen == nil {
+					// No new screen. We're done.
+					a.QueueEvent(nil)
+					return
+				}
 
-			// We have a new screen. Keep going.
-			a.Lock()
-			a.screen = screen
-			enableMouse := a.enableMouse
-			a.Unlock()
+				// We have a new screen. Keep going.
+				a.Lock()
+				a.screen = screen
+				enableMouse := a.enableMouse
+				a.Unlock()
 
-			// Initialize and draw this screen.
-			if err := screen.Init(); err != nil {
-				panic(err)
-			}
-			if enableMouse {
-				screen.EnableMouse()
+				// Initialize and draw this screen.
+				if err := screen.Init(); err != nil {
+					panic(err)
+				}
+				if enableMouse {
+					screen.EnableMouse()
+				}
+				a.draw()
 			}
-			a.draw()
 		}
 	}()
 
 	// Start event loop.
 EventLoop:
-	for {
+	// check to see if the Application.Run is still valid
+	for a.runContext.Err() == nil {
 		select {
-		case event := <-a.events:
-			if event == nil {
+		// break loop when runContext complete
+		case <-a.runContext.Done():
+			break EventLoop
+		case event, ok := <-a.events:
+			if !ok || event == nil {
 				break EventLoop
 			}
 
@@ -348,9 +402,14 @@ EventLoop:
 					if redrawTimer != nil {
 						redrawTimer.Stop()
 					}
-					redrawTimer = time.AfterFunc(redrawPause, func() {
-						a.events <- event
-					})
+					redrawTimer = time.AfterFunc(redrawPause,
+						func() {
+							// check to see if the Application.Run is still valid
+							if a.runContext.Err() == nil {
+								a.events <- event
+							}
+						},
+					)
 				}
 				a.RLock()
 				screen := a.screen
@@ -376,13 +435,19 @@ EventLoop:
 			}
 
 		// If we have updates, now is the time to execute them.
-		case update := <-a.updates:
+		case update, ok := <-a.updates:
+			if !ok {
+				break EventLoop
+			}
 			update.f()
 			if update.done != nil {
 				update.done <- struct{}{}
 			}
 		}
 	}
+	// call the runCancelFunc when exiting eventLoop.
+	//This will stop the channels accepting any more events
+	a.runCancelFunc()
 
 	// Wait for the event loop to finish.
 	wg.Wait()
@@ -500,7 +565,11 @@ func (a *Application) Stop() {
 	}
 	a.screen = nil
 	screen.Fini()
-	a.screenReplacement <- nil
+
+	// check to see if the Application.Run is still valid
+	if a.runContext.Err() == nil {
+		a.screenReplacement <- nil
+	}
 }
 
 // Suspend temporarily suspends the application by exiting terminal UI mode and
@@ -618,15 +687,20 @@ func (a *Application) draw() *Application {
 // corrupted so you may want to offer your users a keyboard shortcut to refresh
 // the screen.
 func (a *Application) Sync() *Application {
-	a.updates <- queuedUpdate{f: func() {
-		a.RLock()
-		screen := a.screen
-		a.RUnlock()
-		if screen == nil {
-			return
+	// check to see if the Application.Run is still valid
+	if a.runContext.Err() == nil {
+		a.updates <- queuedUpdate{
+			f: func() {
+				a.RLock()
+				screen := a.screen
+				a.RUnlock()
+				if screen == nil {
+					return
+				}
+				screen.Sync()
+			},
 		}
-		screen.Sync()
-	}}
+	}
 	return a
 }
 
@@ -741,9 +815,15 @@ func (a *Application) GetFocus() Primitive {
 //
 // This function returns after f has executed.
 func (a *Application) QueueUpdate(f func()) *Application {
-	ch := make(chan struct{})
-	a.updates <- queuedUpdate{f: f, done: ch}
-	<-ch
+	// check to see if the Application.Run is still valid
+	if a.runContext.Err() == nil {
+		ch := make(chan struct{})
+		a.updates <- queuedUpdate{
+			f:    f,
+			done: ch,
+		}
+		<-ch
+	}
 	return a
 }
 
@@ -761,6 +841,9 @@ func (a *Application) QueueUpdateDraw(f func()) *Application {
 //
 // It is not recommended for event to be nil.
 func (a *Application) QueueEvent(event tcell.Event) *Application {
-	a.events <- event
+	// check to see if the Application.Run is still valid
+	if a.runContext.Err() == nil {
+		a.events <- event
+	}
 	return a
 }