orangesite/hn/fetcher.go

package hn

import (
	"net/http"
	"os"
	"time"

	"github.com/jinzhu/gorm"
	// sqlite3 dialect registered with gorm via import side effects.
	_ "github.com/jinzhu/gorm/dialects/sqlite"
	"github.com/peterhellberg/hn"

	"github.com/rs/zerolog"
)

const (
	// frontPageSize is the number of stories shown on the HN front page.
	frontPageSize = 30
	// defaultFetchInterval is how often the front page is scraped.
	defaultFetchInterval = 60 * time.Second
	// httpClientTimeout bounds each request to the HN API.
	httpClientTimeout = 5 * time.Second
)

// Fetcher periodically scrapes the HN front page and records front-page
// tenure in the database.
type Fetcher struct {
	nextFetch     time.Time
	fetchInterval time.Duration
	db            *gorm.DB
	hn            *hn.Client
	log           *zerolog.Logger
}

// NewFetcher builds a Fetcher that scrapes the HN front page once a minute.
func NewFetcher(db *gorm.DB) *Fetcher {
	f := new(Fetcher)
	f.db = db
	f.fetchInterval = defaultFetchInterval
	f.hn = hn.NewClient(&http.Client{
		Timeout: httpClientTimeout,
	})

	return f
}

// AddLogger attaches a logger to the fetcher.
func (f *Fetcher) AddLogger(l *zerolog.Logger) {
	f.log = l
}

// StoreFrontPage scrapes the current HN front page, records new and changed
// stories, and marks stories that have left the front page.
func (f *Fetcher) StoreFrontPage() error {
	ids, err := f.hn.TopStories()
	if err != nil {
		return err
	}

	t := time.Now()

	for i, id := range ids[:frontPageSize] {
		item, err := f.hn.Item(id)
		if err != nil {
			return err
		}

		// check to see if the item was on the frontpage already or not
		var c int
		f.db.Model(&HNFrontPage{}).Where("hn_id = ?", id).Count(&c)

		if c == 0 {
			f.recordNewStory(id, i, t, item)
		} else {
			f.updateExistingStory(id, i, item)
		}
	}

	return f.markDepartedStories(ids, t)
}

// recordNewStory records a story appearing on the front page for the first
// time.
func (f *Fetcher) recordNewStory(id, rank int, t time.Time, item *hn.Item) {
	r := HNFrontPage{
		HNID:        uint(id),
		Appeared:    t,
		Disappeared: time.Time{},
		HighestRank: uint(rank + 1),
		Rank:        uint(rank + 1),
		Title:       item.Title,
		Score:       uint(item.Score),
		URL:         item.URL,
	}
	f.db.Create(&r)
	f.log.Info().
		Uint("hnid", uint(id)).
		Uint("rank", uint(rank+1)).
		Str("title", item.Title).
		Int("score", item.Score).
		Str("url", item.URL).
		Msg("HN new story on frontpage")
}

// updateExistingStory updates the rank, record rank, and score of a story
// still on (or returned to) the front page.
func (f *Fetcher) updateExistingStory(id, rank int, item *hn.Item) {
	var old HNFrontPage
	f.db.Model(&HNFrontPage{}).Where("hn_id = ?", id).First(&old)

	if old.Rank != uint(rank+1) {
		f.log.Info().
			Uint("hnid", uint(id)).
			Uint("oldrank", old.Rank).
			Uint("newrank", uint(rank+1)).
			Int("score", item.Score).
			Str("title", item.Title).
			Str("url", item.URL).
			Msg("HN story rank changed, recording new rank")

		old.Rank = uint(rank + 1)
		old.Score = uint(item.Score)
	}

	if old.HighestRank > uint(rank+1) {
		f.log.Info().
			Uint("hnid", uint(id)).
			Uint("oldrecord", old.HighestRank).
			Uint("newrecord", uint(rank+1)).
			Msg("recording new record high rank for story")

		old.HighestRank = uint(rank + 1)
	}

	if old.Score != uint(item.Score) {
		old.Score = uint(item.Score)
	}

	// in any case it's here now
	old.Disappeared = time.Time{}
	f.db.Save(&old)
}

// markDepartedStories marks any active front-page rows whose story no longer
// appears in the current scrape as departed at time t.
func (f *Fetcher) markDepartedStories(ids []int, t time.Time) error {
	fpitems, err := f.db.Model(&HNFrontPage{}).Where("disappeared is ?", time.Time{}).Rows()
	if err != nil {
		f.log.Error().Err(err).Msg("querying active frontpage items")

		return err
	}

	var toupdate []uint

	for fpitems.Next() {
		var item HNFrontPage

		err = f.db.ScanRows(fpitems, &item)
		if err != nil {
			f.log.Error().Err(err).Msg("scanning frontpage row")

			continue
		}

		exitedFrontPage := true

		for _, xd := range ids[:frontPageSize] {
			if item.HNID == uint(xd) {
				exitedFrontPage = false
			}
		}

		if exitedFrontPage {
			toupdate = append(toupdate, item.HNID)
			dur := t.Sub(item.Appeared).String()
			f.log.Info().
				Uint("hnid", item.HNID).
				Uint("HighestRank", item.HighestRank).
				Str("title", item.Title).
				Str("time_on_frontpage", dur).
				Str("url", item.URL).
				Msg("HN story exited frontpage")
		}
	}

	// close result before we do the update
	_ = fpitems.Close()

	f.db.Model(&HNFrontPage{}).Where("disappeared is ? and hn_id in (?)", time.Time{}, toupdate).Update("Disappeared", t)

	return nil
}

// run migrates the schema and then scrapes the front page forever on the
// configured interval.
func (f *Fetcher) run() {
	if os.Getenv("DEBUG") != "" {
		f.db.LogMode(true)
	}

	f.db.AutoMigrate(&HNStoryRank{})
	f.db.AutoMigrate(&FrontPageCache{})
	f.db.AutoMigrate(&HNFrontPage{})

	for {
		f.log.Info().Msg("fetching top stories from HN")
		f.nextFetch = time.Now().Add(f.fetchInterval)

		err := f.StoreFrontPage()
		if err != nil {
			panic(err)
		}

		until := time.Until(f.nextFetch)
		countdown := time.NewTimer(until)

		f.log.Info().Msgf("waiting %s until next fetch", until)
		<-countdown.C
	}
}