Stop the beat when no more videos to fetch

This commit is contained in:
Gabriel Augendre 2018-11-18 18:46:54 +01:00
parent d0ded4fd02
commit 007dbdd3a9
1 changed files with 18 additions and 9 deletions

View File

@ -35,7 +35,7 @@ func New(b *beat.Beat, cfg *common.Config) (beat.Beater, error) {
return bt, nil return bt, nil
} }
func scrapeVideos(startId string, fieldsToSend chan common.MapStr) { func scrapeVideos(startId string, fieldsToSend chan common.MapStr, done chan bool) {
const BaseUrl = "https://www.youtube.com" const BaseUrl = "https://www.youtube.com"
const BaseSuffix = "/watch?v=" const BaseSuffix = "/watch?v="
@ -81,6 +81,7 @@ func scrapeVideos(startId string, fieldsToSend chan common.MapStr) {
videoCollector.Visit(BaseUrl + BaseSuffix + startId) videoCollector.Visit(BaseUrl + BaseSuffix + startId)
videoCollector.Wait() videoCollector.Wait()
done <- true
} }
// Run starts youtubebeat. // Run starts youtubebeat.
@ -94,7 +95,8 @@ func (bt *Youtubebeat) Run(b *beat.Beat) error {
} }
fieldsToSend := make(chan common.MapStr) fieldsToSend := make(chan common.MapStr)
go scrapeVideos(bt.config.StartId, fieldsToSend) done := make(chan bool)
go scrapeVideos(bt.config.StartId, fieldsToSend, done)
ticker := time.NewTicker(bt.config.Period) ticker := time.NewTicker(bt.config.Period)
for { for {
@ -103,14 +105,21 @@ func (bt *Youtubebeat) Run(b *beat.Beat) error {
return nil return nil
case <-ticker.C: case <-ticker.C:
} }
fields := <-fieldsToSend // Handle a SIGINT even when no more videos to fetch
fields["type"] = b.Info.Name select {
event := beat.Event{ case <-done:
Timestamp: time.Now(), return nil
Fields: fields, case <-bt.done:
return nil
case fields := <-fieldsToSend:
fields["type"] = b.Info.Name
event := beat.Event{
Timestamp: time.Now(),
Fields: fields,
}
bt.client.Publish(event)
logp.Info("Event sent")
} }
bt.client.Publish(event)
logp.Info("Event sent")
} }
} }