<commit_before>package queue

import (
	"os"

	"cloud.google.com/go/pubsub"
	"code.cloudfoundry.org/lager"
	"github.com/tedsuo/ifrit"
	"golang.org/x/net/context"

	"cred-alert/metrics"
)

const SubscriberCount = 2

type pubSubSubscriber struct {
	logger       lager.Logger
	subscription *pubsub.Subscription
	processor    PubSubProcessor

	processTimer          metrics.Timer
	processSuccessCounter metrics.Counter
	processFailureCounter metrics.Counter
	processRetryCounter   metrics.Counter
}

func NewPubSubSubscriber(
	logger lager.Logger,
	subscription *pubsub.Subscription,
	processor PubSubProcessor,
	emitter metrics.Emitter,
) ifrit.Runner {
	return &pubSubSubscriber{
		logger: logger.Session("message-processor", lager.Data{
			"subscription": subscription.ID(),
		}),
		subscription:          subscription,
		processor:             processor,
		processTimer:          emitter.Timer("revok.pub_sub_subscriber.process.time"),
		processSuccessCounter: emitter.Counter("revok.pub_sub_subscriber.process.success"),
		processFailureCounter: emitter.Counter("revok.pub_sub_subscriber.process.failure"),
		processRetryCounter:   emitter.Counter("revok.pub_sub_subscriber.process.retries"),
	}
}

func (p *pubSubSubscriber) Run(signals <-chan os.Signal, ready chan<- struct{}) error {
	p.logger.Info("starting")

	errs := make(chan error)
	cctx, cancel := context.WithCancel(context.Background())

	go func() {
		errs <- p.subscription.Receive(cctx, func(ctx context.Context, message *pubsub.Message) {
			p.processMessage(ctx, message)
		})
	}()

	p.logger.Info("started")

	close(ready)

	select {
	case <-signals:
		p.logger.Info("signalled")
	case err := <-errs:
		if err != nil {
			p.logger.Error("failed", err)
			return err
		}
	}

	cancel()

	p.logger.Info("done")

	return nil
}

func (p *pubSubSubscriber) processMessage(ctx context.Context, message *pubsub.Message) {
	var (
		retryable bool
		err       error
	)

	logger := p.logger.Session("processing-message", lager.Data{
		"pubsub-message":      message.ID,
		"pubsub-publish-time": message.PublishTime.String(),
	})

	p.processTimer.Time(logger, func() {
		retryable, err = p.processor.Process(logger, message)
	})

	if err != nil {
		logger.Error("failed-to-process-message", err)

		if retryable {
			logger.Info("queuing-message-for-retry")
			message.Nack()
			p.processRetryCounter.Inc(logger)
		} else {
			message.Ack()
		}

		p.processFailureCounter.Inc(logger)
	} else {
		message.Ack()
		p.processSuccessCounter.Inc(logger)
	}
}
<commit_msg>limit amount of in-flight fetches to 4<commit_after>package queue

import (
	"os"

	"cloud.google.com/go/pubsub"
	"code.cloudfoundry.org/lager"
	"github.com/tedsuo/ifrit"
	"golang.org/x/net/context"

	"cred-alert/metrics"
)

const SubscriberCount = 2

type pubSubSubscriber struct {
	logger       lager.Logger
	subscription *pubsub.Subscription
	processor    PubSubProcessor

	processTimer          metrics.Timer
	processSuccessCounter metrics.Counter
	processFailureCounter metrics.Counter
	processRetryCounter   metrics.Counter
}

func NewPubSubSubscriber(
	logger lager.Logger,
	subscription *pubsub.Subscription,
	processor PubSubProcessor,
	emitter metrics.Emitter,
) ifrit.Runner {
	return &pubSubSubscriber{
		logger: logger.Session("message-processor", lager.Data{
			"subscription": subscription.ID(),
		}),
		subscription:          subscription,
		processor:             processor,
		processTimer:          emitter.Timer("revok.pub_sub_subscriber.process.time"),
		processSuccessCounter: emitter.Counter("revok.pub_sub_subscriber.process.success"),
		processFailureCounter: emitter.Counter("revok.pub_sub_subscriber.process.failure"),
		processRetryCounter:   emitter.Counter("revok.pub_sub_subscriber.process.retries"),
	}
}

func (p *pubSubSubscriber) Run(signals <-chan os.Signal, ready chan<- struct{}) error {
	p.logger.Info("starting")

	errs := make(chan error)
	cctx, cancel := context.WithCancel(context.Background())

	p.subscription.MaxOutstandingMessages = 4

	go func() {
		errs <- p.subscription.Receive(cctx, func(ctx context.Context, message *pubsub.Message) {
			p.processMessage(ctx, message)
		})
	}()

	p.logger.Info("started")

	close(ready)

	select {
	case <-signals:
		p.logger.Info("signalled")
	case err := <-errs:
		if err != nil {
			p.logger.Error("failed", err)
			return err
		}
	}

	cancel()

	p.logger.Info("done")

	return nil
}

func (p *pubSubSubscriber) processMessage(ctx context.Context, message *pubsub.Message) {
	var (
		retryable bool
		err       error
	)

	logger := p.logger.Session("processing-message", lager.Data{
		"pubsub-message":      message.ID,
		"pubsub-publish-time": message.PublishTime.String(),
	})

	p.processTimer.Time(logger, func() {
		retryable, err = p.processor.Process(logger, message)
	})

	if err != nil {
		logger.Error("failed-to-process-message", err)

		if retryable {
			logger.Info("queuing-message-for-retry")
			message.Nack()
			p.processRetryCounter.Inc(logger)
		} else {
			message.Ack()
		}

		p.processFailureCounter.Inc(logger)
	} else {
		message.Ack()
		p.processSuccessCounter.Inc(logger)
	}
}

{'fraction_non_alphanumeric': 0.12986504987287306, 'fraction_numerical': 0.0007823195775474281, 'mean_word_length': 4.108891108891109, 'pattern_counts': {'":': 6, '<': 13, '<?xml version=': 0, '>': 3, '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': '2482601', 'n_tokens_mistral': 1832, 'n_tokens_neox': 1704, 'n_words': 302}