Skip to content

Commit

Permalink
bot,api: fix stop channel race
Browse files Browse the repository at this point in the history
  • Loading branch information
davseby authored and demget committed Feb 20, 2024
1 parent 71ac299 commit 4cce303
Show file tree
Hide file tree
Showing 2 changed files with 20 additions and 7 deletions.
13 changes: 8 additions & 5 deletions api.go
Original file line number Diff line number Diff line change
Expand Up @@ -27,18 +27,21 @@ func (b *Bot) Raw(method string, payload interface{}) ([]byte, error) {
return nil, err
}

// Cancel the request immediately without waiting for the timeout when bot is about to stop.
// Cancel the request immediately without waiting for the timeout
// when bot is about to stop.
// This may become important if doing long polling with long timeout.
exit := make(chan struct{})
defer close(exit)
ctx, cancel := context.WithCancel(context.Background())
defer cancel()

go func() {
b.stopMu.RLock()
stopCh := b.stopClient
b.stopMu.RUnlock()

select {
case <-b.stopClient:
case <-stopCh:
cancel()
case <-exit:
case <-ctx.Done():
}
}()

Expand Down
14 changes: 12 additions & 2 deletions bot.go
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ import (
"regexp"
"strconv"
"strings"
"sync"
"time"
)

Expand Down Expand Up @@ -81,7 +82,9 @@ type Bot struct {
parseMode ParseMode
stop chan chan struct{}
client *http.Client
stopClient chan struct{}

stopMu sync.RWMutex
stopClient chan struct{}
}

// Settings represents a utility struct for passing certain
Expand Down Expand Up @@ -198,10 +201,14 @@ func (b *Bot) Start() {
}

// do nothing if called twice
b.stopMu.Lock()
if b.stopClient != nil {
b.stopMu.Unlock()
return
}

b.stopClient = make(chan struct{})
b.stopMu.Unlock()

stop := make(chan struct{})
stopConfirm := make(chan struct{})
Expand All @@ -221,17 +228,20 @@ func (b *Bot) Start() {
close(stop)
<-stopConfirm
close(confirm)
b.stopClient = nil
return
}
}
}

// Stop gracefully shuts the poller down.
func (b *Bot) Stop() {
b.stopMu.Lock()
if b.stopClient != nil {
close(b.stopClient)
b.stopClient = nil
}
b.stopMu.Unlock()

confirm := make(chan struct{})
b.stop <- confirm
<-confirm
Expand Down

0 comments on commit 4cce303

Please sign in to comment.