Fixed race condition

mymmrac · May 30, 2022 · c1109f4 · c1109f4
1 parent 35cddeb
commit c1109f4
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 0 deletions.
diff --git a/telegohandler/bot_handler.go b/telegohandler/bot_handler.go
@@ -60,9 +60,13 @@ func (h *BotHandler) Start() {
 	h.running = true
 	h.runningLock.Unlock()
 
+	// Prevents calling Wait before single Add call
+	h.handledUpdates.Add(1)
+
 	for {
 		select {
 		case <-h.stop:
+			h.handledUpdates.Done()
 			return
 		case update := <-h.updates:
 			h.processUpdate(update)

diff --git a/telegohandler/bot_handler_test.go b/telegohandler/bot_handler_test.go
@@ -151,6 +151,9 @@ func TestBotHandler_Stop(t *testing.T) {
 
 		assert.NotPanics(t, func() {
 			go bh.Start()
+			for !bh.IsRunning() {
+				// Wait for handler to start
+			}
 
 			updates <- telego.Update{}
 
@@ -183,6 +186,9 @@ func TestBotHandler_Stop(t *testing.T) {
 
 		assert.NotPanics(t, func() {
 			go bh.Start()
+			for !bh.IsRunning() {
+				// Wait for handler to start
+			}
 
 			updates <- telego.Update{}
 
@@ -198,6 +204,31 @@ func TestBotHandler_Stop(t *testing.T) {
 			}
 		})
 	})
+
+	t.Run("stop_checked", func(t *testing.T) {
+		bot, err := telego.NewBot(token)
+		require.NoError(t, err)
+
+		updates := make(chan telego.Update, 1)
+
+		bh, err := NewBotHandler(bot, updates)
+		require.NoError(t, err)
+
+		bh.Handle(func(bot *telego.Bot, update telego.Update) {
+			t.Fatal("handled after stop")
+		})
+
+		assert.NotPanics(t, func() {
+			go bh.Start()
+			for !bh.IsRunning() {
+				// Wait for handler to start
+			}
+
+			bh.Stop()
+
+			updates <- telego.Update{}
+		})
+	})
 }
 
 func TestBotHandler_Handle(t *testing.T) {