Modified poller API, fixes #106.

7 years ago · 119a4e1698
parent 3329d98249
commit 119a4e1698
2 changed files with 148 additions and 121 deletions
--- a/bot.go
+++ b/bot.go
@ -26,6 +26,7 @@ func NewBot(pref Settings) (*Bot, error) {
 		Poller:  pref.Poller,

 		handlers: make(map[string]interface{}),
+		stop:     make(chan struct{}),
 	}

 	user, err := bot.getMe()
@ -124,150 +125,167 @@ func (b *Bot) Start() {
 		panic("telebot: can't start without a poller")
 	}

-	b.stop = make(chan struct{})
-	go b.Poller.Poll(b, b.Updates, b.stop)
+	stopPoller := make(chan struct{})

-	for upd := range b.Updates {
-		if upd.Message != nil {
-			m := upd.Message
+	go b.Poller.Poll(b, b.Updates, stopPoller)

-			if m.PinnedMessage != nil {
-				b.handle(OnPinned, m)
-				continue
-			}
+	for {
+		select {
+		// handle incoming updates
+		case upd := <-b.Updates:
+			b.incomingUpdate(&upd)

-			// Commands
-			if m.Text != "" {
-				// Filtering malicious messsages
-				if m.Text[0] == '\a' {
-					continue
-				}
+		// call to stop polling
+		case <-b.stop:
+			stopPoller <- struct{}{}

-				match := cmdRx.FindAllStringSubmatch(m.Text, -1)
+		// polling has stopped
+		case <-stopPoller:
+			return
+		}
+	}
+}

-				// Command found - handle and continue
-				if match != nil && b.handleCommand(m, match[0][1], match[0][3]) {
-					continue
-				}
+func (b *Bot) incomingUpdate(upd *Update) {
+	if upd.Message != nil {
+		m := upd.Message

-				// OnText
-				b.handle(OnText, m)
-				continue
-			}
+		if m.PinnedMessage != nil {
+			b.handle(OnPinned, m)
+			return
+		}

-			// on media
-			if b.handleMedia(m) {
-				continue
+		// Commands
+		if m.Text != "" {
+			// Filtering malicious messsages
+			if m.Text[0] == '\a' {
+				return
 			}

-			// OnAddedToGrop
-			wasAdded := m.UserJoined.ID == b.Me.ID ||
-				m.UsersJoined != nil && isUserInList(b.Me, m.UsersJoined)
-			if m.GroupCreated || m.SuperGroupCreated || wasAdded {
-				b.handle(OnAddedToGroup, m)
-				continue
-			}
+			match := cmdRx.FindAllStringSubmatch(m.Text, -1)

-			if m.UserJoined != nil {
-				b.handle(OnUserJoined, m)
-				continue
+			// Command found - handle and return
+			if match != nil && b.handleCommand(m, match[0][1], match[0][3]) {
+				return
 			}

-			if m.UsersJoined != nil {
-				for _, user := range m.UsersJoined {
-					m.UserJoined = &user
-					b.handle(OnUserJoined, m)
-				}
-
-				continue
-			}
+			// OnText
+			b.handle(OnText, m)
+			return
+		}

-			if m.UserLeft != nil {
-				b.handle(OnUserLeft, m)
-				continue
-			}
+		// on media
+		if b.handleMedia(m) {
+			return
+		}

-			if m.NewGroupTitle != "" {
-				b.handle(OnNewGroupTitle, m)
-				continue
-			}
+		// OnAddedToGrop
+		wasAdded := m.UserJoined.ID == b.Me.ID ||
+			m.UsersJoined != nil && isUserInList(b.Me, m.UsersJoined)
+		if m.GroupCreated || m.SuperGroupCreated || wasAdded {
+			b.handle(OnAddedToGroup, m)
+			return
+		}

-			if m.NewGroupPhoto != nil {
-				b.handle(OnNewGroupPhoto, m)
-				continue
-			}
+		if m.UserJoined != nil {
+			b.handle(OnUserJoined, m)
+			return
+		}

-			if m.GroupPhotoDeleted {
-				b.handle(OnGroupPhotoDeleted, m)
-				continue
+		if m.UsersJoined != nil {
+			for _, user := range m.UsersJoined {
+				m.UserJoined = &user
+				b.handle(OnUserJoined, m)
 			}

-			if m.MigrateTo != 0 {
-				if handler, ok := b.handlers[OnMigration]; ok {
-					if handler, ok := handler.(func(int64, int64)); ok {
-						handler(m.MigrateFrom, m.MigrateTo)
-					}
-				}
+			return
+		}

-				continue
-			}
+		if m.UserLeft != nil {
+			b.handle(OnUserLeft, m)
+			return
+		}

-			continue
+		if m.NewGroupTitle != "" {
+			b.handle(OnNewGroupTitle, m)
+			return
 		}

-		if upd.EditedMessage != nil {
-			b.handle(OnEdited, upd.EditedMessage)
-			continue
+		if m.NewGroupPhoto != nil {
+			b.handle(OnNewGroupPhoto, m)
+			return
 		}

-		if upd.ChannelPost != nil {
-			b.handle(OnChannelPost, upd.ChannelPost)
-			continue
+		if m.GroupPhotoDeleted {
+			b.handle(OnGroupPhotoDeleted, m)
+			return
 		}

-		if upd.EditedChannelPost != nil {
-			b.handle(OnEditedChannelPost, upd.EditedChannelPost)
-			continue
+		if m.MigrateTo != 0 {
+			if handler, ok := b.handlers[OnMigration]; ok {
+				if handler, ok := handler.(func(int64, int64)); ok {
+					handler(m.MigrateFrom, m.MigrateTo)
+				}
+			}
+
+			return
 		}

-		if upd.Callback != nil {
-			if upd.Callback.Data != "" {
-				data := upd.Callback.Data
+		return
+	}

-				if data[0] == '\f' {
-					match := cbackRx.FindAllStringSubmatch(data, -1)
+	if upd.EditedMessage != nil {
+		b.handle(OnEdited, upd.EditedMessage)
+		return
+	}

-					if match != nil {
-						unique, payload := match[0][1], match[0][2]
+	if upd.ChannelPost != nil {
+		b.handle(OnChannelPost, upd.ChannelPost)
+		return
+	}

-						if handler, ok := b.handlers["\f"+unique]; ok {
-							if handler, ok := handler.(func(*Callback)); ok {
-								upd.Callback.Data = payload
-								handler(upd.Callback)
-								continue
-							}
-						}
+	if upd.EditedChannelPost != nil {
+		b.handle(OnEditedChannelPost, upd.EditedChannelPost)
+		return
+	}
+
+	if upd.Callback != nil {
+		if upd.Callback.Data != "" {
+			data := upd.Callback.Data
+
+			if data[0] == '\f' {
+				match := cbackRx.FindAllStringSubmatch(data, -1)

+				if match != nil {
+					unique, payload := match[0][1], match[0][2]
+
+					if handler, ok := b.handlers["\f"+unique]; ok {
+						if handler, ok := handler.(func(*Callback)); ok {
+							upd.Callback.Data = payload
+							handler(upd.Callback)
+							return
+						}
 					}
+
 				}
 			}
+		}

-			if handler, ok := b.handlers[OnCallback]; ok {
-				if handler, ok := handler.(func(*Callback)); ok {
-					handler(upd.Callback)
-				}
+		if handler, ok := b.handlers[OnCallback]; ok {
+			if handler, ok := handler.(func(*Callback)); ok {
+				handler(upd.Callback)
 			}
-			continue
 		}
+		return
+	}

-		if upd.Query != nil {
-			if handler, ok := b.handlers[OnQuery]; ok {
-				if handler, ok := handler.(func(*Query)); ok {
-					handler(upd.Query)
-				}
+	if upd.Query != nil {
+		if handler, ok := b.handlers[OnQuery]; ok {
+			if handler, ok := handler.(func(*Query)); ok {
+				handler(upd.Query)
 			}
-			continue
 		}
+		return
 	}
 }

@ -336,7 +354,7 @@ func (b *Bot) handleMedia(m *Message) bool {

 // Stop gracefully shuts the poller down.
 func (b *Bot) Stop() {
-	close(b.stop)
+	b.stop <- struct{}{}
 }

 // Send accepts 2+ arguments, starting with destination chat, followed by
--- a/poller.go
+++ b/poller.go
@ -15,6 +15,9 @@ type Poller interface {
 	// Poll is supposed to take the bot object
 	// subscription channel and start polling
 	// for Updates immediately.
+	//
+	// Poller must listen for stop constantly and close
+	// it as soon as it's done polling.
 	Poll(b *Bot, updates chan Update, stop chan struct{})
 }

@ -47,15 +50,21 @@ func (p *MiddlewarePoller) Poll(b *Bot, dest chan Update, stop chan struct{}) {
 	}

 	middle := make(chan Update, cap)
-	stop2 := make(chan struct{})
+	stopPoller := make(chan struct{})

-	go p.Poller.Poll(b, middle, stop2)
+	go p.Poller.Poll(b, middle, stopPoller)

 	for {
 		select {
+		// call to stop
 		case <-stop:
-			close(stop2)
+			stopPoller <- struct{}{}
+
+		// poller is done
+		case <-stopPoller:
+			close(stop)
 			return
+
 		case upd := <-middle:
 			if p.filter(&upd) {
 				dest <- upd
@ -67,28 +76,28 @@ func (p *MiddlewarePoller) Poll(b *Bot, dest chan Update, stop chan struct{}) {
 // LongPoller is a classic LongPoller with timeout.
 type LongPoller struct {
 	Timeout time.Duration
+
+	LastUpdateID int
 }

 // Poll does long polling.
 func (p *LongPoller) Poll(b *Bot, dest chan Update, stop chan struct{}) {
-	var latestUpd int
+	go func(stop chan struct{}) {
+		<-stop
+		close(stop)
+	}(stop)

 	for {
-		select {
-		case <-stop:
-			return
-		default:
-			updates, err := b.getUpdates(latestUpd+1, p.Timeout)
+		updates, err := b.getUpdates(p.LastUpdateID+1, p.Timeout)

-			if err != nil {
-				b.debug(errors.Wrap(err, "getUpdates() failed"))
-				continue
-			}
+		if err != nil {
+			b.debug(errors.Wrap(err, "getUpdates() failed"))
+			continue
+		}

-			for _, update := range updates {
-				latestUpd = update.ID
-				dest <- update
-			}
+		for _, update := range updates {
+			p.LastUpdateID = update.ID
+			dest <- update
 		}
 	}
 }