<commit_before>package main

import (
	"io"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/juju/ratelimit"
	"github.com/pkg/errors"
	"github.com/tomnomnom/linkheader"
)

func PubSubMakeCallbackURL(baseURL, id string) string {
	return baseURL + "/" + id
}

type PubSubSubscription struct {
	ID          string
	Hub         string
	Topic       string
	CallbackURL string
	CreatedAt   time.Time
	UpdatedAt   time.Time
	ExpiresAt   *time.Time
}

type PubSubState interface {
	All() (subscriptions []PubSubSubscription, err error)
	Add(hub, topic, baseURL string) (subscription *PubSubSubscription, oldCallbackURL string, err error)
	Get(hub, topic string) (subscription *PubSubSubscription, err error)
	GetByID(id string) (subscription *PubSubSubscription, err error)
	Set(hub, topic string, updatedAt, expiresAt time.Time) (err error)
	Del(hub, topic string) (err error)
}

type PubSubMessageHandler func(id string, s *PubSubSubscription, rd io.ReadCloser)

type PubSubClient struct {
	CallbackURL string
	State       PubSubState
	OnMessage   PubSubMessageHandler
}

func NewPubSubClient(callbackURL string, state PubSubState, onMessage PubSubMessageHandler) *PubSubClient {
	return &PubSubClient{
		CallbackURL: callbackURL,
		State:       state,
		OnMessage:   onMessage,
	}
}

func (c *PubSubClient) Refresh(forceUpdate bool, interval time.Duration) error {
	a, err := c.State.All()
	if err != nil {
		return errors.Wrap(err, "PubSubClient.Refresh")
	}

	logrus.WithField("count", len(a)).Debug("pubsub: got subscriptions to refresh")

	m := make(map[string]*ratelimit.Bucket)
	var l sync.RWMutex

	getBucket := func(host string) *ratelimit.Bucket {
		l.RLock()
		if b, ok := m[host]; ok {
			l.RUnlock()
			return b
		}
		l.RUnlock()

		l.Lock()
		if b, ok := m[host]; ok {
			l.Unlock()
			return b
		}
		defer l.Unlock()

		m[host] = ratelimit.NewBucket(time.Second*30, 4)

		return m[host]
	}

	var g WorkerGroup

	for _, e := range a {
		e := e

		callbackURL := c.CallbackURL + "/" + e.ID

		l := logrus.WithFields(logrus.Fields{
			"id":               e.ID,
			"hub":              e.Hub,
			"topic":            e.Topic,
			"callback_url":     e.CallbackURL,
			"new_callback_url": callbackURL,
			"created_at":       e.CreatedAt,
			"updated_at":       e.UpdatedAt,
			"expires_at":       e.ExpiresAt,
			"force_update":     forceUpdate,
		})

		if e.ExpiresAt == nil {
			l.Debug("pubsub: not refreshing subscription which has never been verified")
			continue
		}

		if forceUpdate || e.ExpiresAt.Sub(time.Now()) < interval || e.CallbackURL != callbackURL {
			l.Debug("pubsub: refreshing subscription")

			g.Add(func() error {
				u, err := url.Parse(e.Hub)
				if err != nil {
					l.WithError(err).Warn("pubsub: couldn't parse hub url")
					return errors.Wrap(err, "PubSubClient.RefreshWorker")
				}

				if dur, skip := getBucket(u.Host).TakeMaxDuration(1, *pubsubRefreshInterval); skip {
					l.Debug("pubsub: skipping renewing for now as we'd have to wait too long")
					return nil
				} else if dur > 0 {
					l.WithField("duration", dur).Debug("pubsub: waiting so as not to overwhelm the endpoint")
					time.Sleep(dur)
				}

				if err := c.Subscribe(e.Hub, e.Topic); err != nil {
					l.WithError(err).Warn("pubsub: couldn't subscribe to topic")<commit_msg>only log when a refresh is actually starting<commit_after>package main

import (
	"io"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/juju/ratelimit"
	"github.com/pkg/errors"
	"github.com/tomnomnom/linkheader"
)

func PubSubMakeCallbackURL(baseURL, id string) string {
	return baseURL + "/" + id
}

type PubSubSubscription struct {
	ID          string
	Hub         string
	Topic       string
	CallbackURL string
	CreatedAt   time.Time
	UpdatedAt   time.Time
	ExpiresAt   *time.Time
}

type PubSubState interface {
	All() (subscriptions []PubSubSubscription, err error)
	Add(hub, topic, baseURL string) (subscription *PubSubSubscription, oldCallbackURL string, err error)
	Get(hub, topic string) (subscription *PubSubSubscription, err error)
	GetByID(id string) (subscription *PubSubSubscription, err error)
	Set(hub, topic string, updatedAt, expiresAt time.Time) (err error)
	Del(hub, topic string) (err error)
}

type PubSubMessageHandler func(id string, s *PubSubSubscription, rd io.ReadCloser)

type PubSubClient struct {
	CallbackURL string
	State       PubSubState
	OnMessage   PubSubMessageHandler
}

func NewPubSubClient(callbackURL string, state PubSubState, onMessage PubSubMessageHandler) *PubSubClient {
	return &PubSubClient{
		CallbackURL: callbackURL,
		State:       state,
		OnMessage:   onMessage,
	}
}

func (c *PubSubClient) Refresh(forceUpdate bool, interval time.Duration) error {
	a, err := c.State.All()
	if err != nil {
		return errors.Wrap(err, "PubSubClient.Refresh")
	}

	logrus.WithField("count", len(a)).Debug("pubsub: got subscriptions to refresh")

	m := make(map[string]*ratelimit.Bucket)
	var l sync.RWMutex

	getBucket := func(host string) *ratelimit.Bucket {
		l.RLock()
		if b, ok := m[host]; ok {
			l.RUnlock()
			return b
		}
		l.RUnlock()

		l.Lock()
		if b, ok := m[host]; ok {
			l.Unlock()
			return b
		}
		defer l.Unlock()

		m[host] = ratelimit.NewBucket(time.Second*30, 4)

		return m[host]
	}

	var g WorkerGroup

	for _, e := range a {
		e := e

		callbackURL := c.CallbackURL + "/" + e.ID

		l := logrus.WithFields(logrus.Fields{
			"id":               e.ID,
			"hub":              e.Hub,
			"topic":            e.Topic,
			"callback_url":     e.CallbackURL,
			"new_callback_url": callbackURL,
			"created_at":       e.CreatedAt,
			"updated_at":       e.UpdatedAt,
			"expires_at":       e.ExpiresAt,
			"force_update":     forceUpdate,
		})

		if e.ExpiresAt == nil {
			l.Debug("pubsub: not refreshing subscription which has never been verified")
			continue
		}

		if forceUpdate || e.ExpiresAt.Sub(time.Now()) < interval || e.CallbackURL != callbackURL {
			g.Add(func() error {
				u, err := url.Parse(e.Hub)
				if err != nil {
					l.WithError(err).Warn("pubsub: couldn't parse hub url")
					return errors.Wrap(err, "PubSubClient.RefreshWorker")
				}

				if dur, skip := getBucket(u.Host).TakeMaxDuration(1, *pubsubRefreshInterval); skip {
					l.Debug("pubsub: skipping renewing for now as we'd have to wait too long")
					return nil
				} else if dur > 0 {
					l.WithField("duration", dur).Debug("pubsub: waiting so as not to overwhelm the endpoint")
					time.Sleep(dur)
				}

				l.Debug("pubsub: refreshing subscription")

				if err := c.Subscribe(e.Hub, e.Topic); err != nil {
					l.WithError(err).Warn("pubsub: couldn't subscribe to topic")
{'fraction_non_alphanumeric': 0.12518628912071536, 'fraction_numerical': 0.0014903129657228018, 'mean_word_length': 3.406434668417597, 'pattern_counts': {'":': 18, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1236573', 'n_tokens_mistral': 2545, 'n_tokens_neox': 2299, 'n_words': 586}