<commit_before>// Package gcssource implements files.Source on top of Google Cloud Storage.
package gcssource

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"

	"cloud.google.com/go/pubsub"
	"cloud.google.com/go/storage"
	"go.skia.org/infra/go/auth"
	"go.skia.org/infra/go/httputils"
	"go.skia.org/infra/go/metrics2"
	"go.skia.org/infra/go/pubsub/sub"
	"go.skia.org/infra/go/skerr"
	"go.skia.org/infra/go/sklog"
	"go.skia.org/infra/perf/go/config"
	"go.skia.org/infra/perf/go/file"
	"go.skia.org/infra/perf/go/ingest/filter"
	"google.golang.org/api/option"
)

const (
	// maxParallelReceives is the number of Go routines we want to run.
	// Determined experimentally.
	maxParallelReceives = 1

	// subscriptionSuffix is the name we append to a topic name to build a
	// subscription name.
	subscriptionSuffix = "-prod"
)

// pubSubEvent is used to deserialize the PubSub data.
//
// The PubSub event data is a JSON serialized storage.ObjectAttrs object.
// See https://cloud.google.com/storage/docs/pubsub-notifications#payload
type pubSubEvent struct {
	Bucket string `json:"bucket"`
	Name   string `json:"name"`
}

// GCSSource implements file.Source for Google Cloud Storage.
type GCSSource struct {
	// instanceConfig if the InstanceConfig we are ingesting files for.
	instanceConfig *config.InstanceConfig

	// storageClient is how we talk to Google Cloud Storage.
	storageClient *storage.Client

	// fileChannel is the output channel returned from Start.
	fileChannel chan<- file.File

	// subscription is the pubsub event subscription.
	subscription *pubsub.Subscription

	// started is true if Start has already been called.
	started bool

	// nackCounter is a metric of how many messages we've nacked.
	nackCounter metrics2.Counter

	// ackCounter is a metric of how many messages we've acked.
	ackCounter metrics2.Counter

	// filter to accept/reject files based on their filename.
	filter *filter.Filter
}

// New returns a new *GCSSource
func New(ctx context.Context, instanceConfig *config.InstanceConfig, local bool) (*GCSSource, error) {
	ts, err := auth.NewDefaultTokenSource(local, storage.ScopeReadOnly, pubsub.ScopePubSub)
	if err != nil {
		return nil, skerr.Wrap(err)
	}
	client := httputils.DefaultClientConfig().WithTokenSource(ts).WithoutRetries().Client()
	gcsClient, err := storage.NewClient(ctx, option.WithHTTPClient(client))
	if err != nil {
		return nil, skerr.Wrap(err)
	}

	sub, err := sub.New(ctx, local, instanceConfig.IngestionConfig.SourceConfig.Project, instanceConfig.IngestionConfig.SourceConfig.Topic, maxParallelReceives)
	if err != nil {
		return nil, skerr.Wrap(err)
	}

	f, err := filter.New(instanceConfig.IngestionConfig.SourceConfig.AcceptIfNameMatches, instanceConfig.IngestionConfig.SourceConfig.RejectIfNameMatches)
	if err != nil {
		return nil, skerr.Wrap(err)
	}<commit_msg>[perf] gcssource uses SourceConfig.Subscription.<commit_after>// Package gcssource implements files.Source on top of Google Cloud Storage.
package gcssource

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"

	"cloud.google.com/go/pubsub"
	"cloud.google.com/go/storage"
	"go.skia.org/infra/go/auth"
	"go.skia.org/infra/go/httputils"
	"go.skia.org/infra/go/metrics2"
	"go.skia.org/infra/go/pubsub/sub"
	"go.skia.org/infra/go/skerr"
	"go.skia.org/infra/go/sklog"
	"go.skia.org/infra/perf/go/config"
	"go.skia.org/infra/perf/go/file"
	"go.skia.org/infra/perf/go/ingest/filter"
	"google.golang.org/api/option"
)

const (
	// maxParallelReceives is the number of Go routines we want to run.
	// Determined experimentally.
	maxParallelReceives = 1

	// subscriptionSuffix is the name we append to a topic name to build a
	// subscription name.
	subscriptionSuffix = "-prod"
)

// pubSubEvent is used to deserialize the PubSub data.
//
// The PubSub event data is a JSON serialized storage.ObjectAttrs object.
// See https://cloud.google.com/storage/docs/pubsub-notifications#payload
type pubSubEvent struct {
	Bucket string `json:"bucket"`
	Name   string `json:"name"`
}

// GCSSource implements file.Source for Google Cloud Storage.
type GCSSource struct {
	// instanceConfig if the InstanceConfig we are ingesting files for.
	instanceConfig *config.InstanceConfig

	// storageClient is how we talk to Google Cloud Storage.
	storageClient *storage.Client

	// fileChannel is the output channel returned from Start.
	fileChannel chan<- file.File

	// subscription is the pubsub event subscription.
	subscription *pubsub.Subscription

	// started is true if Start has already been called.
	started bool

	// nackCounter is a metric of how many messages we've nacked.
	nackCounter metrics2.Counter

	// ackCounter is a metric of how many messages we've acked.
	ackCounter metrics2.Counter

	// filter to accept/reject files based on their filename.
	filter *filter.Filter
}

// New returns a new *GCSSource
func New(ctx context.Context, instanceConfig *config.InstanceConfig, local bool) (*GCSSource, error) {
	ts, err := auth.NewDefaultTokenSource(local, storage.ScopeReadOnly, pubsub.ScopePubSub)
	if err != nil {
		return nil, skerr.Wrap(err)
	}
	client := httputils.DefaultClientConfig().WithTokenSource(ts).WithoutRetries().Client()
	gcsClient, err := storage.NewClient(ctx, option.WithHTTPClient(client))
	if err != nil {
		return nil, skerr.Wrap(err)
	}

	subName := instanceConfig.IngestionConfig.SourceConfig.Subscription
	if subName == "" {
		subName, err = sub.NewRoundRobinNameProvider(local, instanceConfig.IngestionConfig.SourceConfig.Topic).SubName()
		if err != nil {
			return nil, skerr.Wrap(err)
		}
	}
	sub, err := sub.NewWithSubName(ctx, local, instanceConfig.IngestionConfig.SourceConfig.Project, instanceConfig.IngestionConfig.SourceConfig.Topic, subName, maxParallelReceives)
	if err != nil {
		return nil, skerr.Wrap(err)
	}

	f, err := filter.New(instanceConfig.IngestionConfig.SourceConfig.AcceptIfNameMatches, instanceConfig.IngestionConfig.SourceConfig.RejectIfNameMatches)
	if err != nil {
		return nil, skerr.Wrap(err)
	}
{'fraction_non_alphanumeric': 0.10997827177001504, 'fraction_numerical': 0.0013371218452281465, 'mean_word_length': 4.309671694764862, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 3, 'https://': 2, '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': '3293283', 'n_tokens_mistral': 2006, 'n_tokens_neox': 1944, 'n_words': 514}