| package ingestion |
| |
| import ( |
| "bytes" |
| "context" |
| "encoding/hex" |
| "errors" |
| "fmt" |
| "io" |
| "io/ioutil" |
| "net/http" |
| "regexp" |
| "sync" |
| "time" |
| |
| "cloud.google.com/go/storage" |
| "github.com/cenkalti/backoff" |
| "go.skia.org/infra/go/eventbus" |
| "go.skia.org/infra/go/fileutil" |
| "go.skia.org/infra/go/gcs" |
| "go.skia.org/infra/go/git/gitinfo" |
| "go.skia.org/infra/go/gitiles" |
| "go.skia.org/infra/go/gitstore/bt_gitstore" |
| "go.skia.org/infra/go/sharedconfig" |
| "go.skia.org/infra/go/skerr" |
| "go.skia.org/infra/go/sklog" |
| "go.skia.org/infra/go/util" |
| "go.skia.org/infra/go/vcsinfo" |
| "go.skia.org/infra/go/vcsinfo/bt_vcs" |
| "google.golang.org/api/option" |
| ) |
| |
| const ( |
| // Limit the ingester's attempts to get a file before giving up. |
| maxReadTime = 30 * time.Second |
| |
| // maxConcurrentDirPollers is the maximum number of concurrent go-routines that |
| // read from GCS and the file system when polling a range of directories. |
| maxConcurrentDirPollers = 200 |
| ) |
| |
| // Constructor is the signature that has to be implemented to register a |
| // Processor implementation to be instantiated by name from a config struct. |
| // vcs is an instance that might be shared across multiple ingesters. |
| // config is usually parsed from a JSON5 file. |
| // client can be assumed to be ready to serve the needs of the resulting Processor. |
| // eventBus is the eventbus to be used by the ingester (optional). |
| type Constructor func(context.Context, vcsinfo.VCS, *sharedconfig.IngesterConfig, *http.Client) (Processor, error) |
| |
| // stores the constructors that register for instantiation from a config struct. |
| var constructors = map[string]Constructor{} |
| |
| // used to synchronize constructor registration and instantiation. |
| var registrationMutex sync.Mutex |
| |
| // Register registers the given constructor to create an instance of a Processor. |
| func Register(id string, constructor Constructor) { |
| registrationMutex.Lock() |
| defer registrationMutex.Unlock() |
| constructors[id] = constructor |
| } |
| |
| // IngestersFromConfig creates a list of ingesters from a config struct. |
| // Usually the struct is created from parsing a config file. |
| // client is assumed to be suitable for the given application. If e.g. the |
| // processors of the current application require an authenticated http client, |
| // then it is expected that client meets these requirements. |
| func IngestersFromConfig(ctx context.Context, config *sharedconfig.Config, client *http.Client, eventBus eventbus.EventBus, ingestionStore IngestionStore, btConf *bt_gitstore.BTConfig) ([]*Ingester, error) { |
| if client == nil { |
| return nil, errors.New("httpClient cannot be nil") |
| } |
| |
| registrationMutex.Lock() |
| defer registrationMutex.Unlock() |
| ret := []*Ingester{} |
| |
| // Make sure we have an eventbus since that is shared by ingesters and sources. |
| if eventBus == nil { |
| eventBus = eventbus.New() |
| } |
| |
| // Set up the gitinfo object. |
| var vcs vcsinfo.VCS |
| if btConf != nil { |
| gitStore, err := bt_gitstore.New(ctx, btConf, config.GitRepoURL) |
| if err != nil { |
| return nil, skerr.Wrapf(err, "could not instantiate gitstore for %s", config.GitRepoURL) |
| } |
| |
| // Set up VCS instance to track master. |
| gitilesRepo := gitiles.NewRepo(config.GitRepoURL, client) |
| if vcs, err = bt_vcs.New(ctx, gitStore, "master", gitilesRepo); err != nil { |
| return nil, skerr.Wrapf(err, "could not create new bt_vcs") |
| } |
| |
| sklog.Infof("Created vcs client based on BigTable.") |
| } else { |
| var err error |
| if vcs, err = gitinfo.CloneOrUpdate(ctx, config.GitRepoURL, config.GitRepoDir, true); err != nil { |
| return nil, skerr.Wrapf(err, "could not clone %s locally to %s", config.GitRepoURL, config.GitRepoDir) |
| } |
| sklog.Infof("Created vcs client based on local checkout.") |
| } |
| |
| // Instantiate the secondary repo if one was specified. |
| // TODO(kjlubick): make this support bigtable git also. |
| if config.SecondaryRepoURL != "" { |
| // TODO(kjlubick) Check up tracestore_impl's isOnMaster to make sure it |
| // works with what is put here. |
| return nil, skerr.Fmt("Not yet implemented to have a secondary repo url") |
| } |
| |
| // for each defined ingester create an instance. |
| for id, ingesterConf := range config.Ingesters { |
| sklog.Infof("Starting to instantiate ingester: %s", id) |
| processorConstructor, ok := constructors[id] |
| if !ok { |
| return nil, skerr.Fmt("unknown ingester: '%s'", id) |
| } |
| |
| // Instantiate the sources |
| sources := make([]Source, 0, len(ingesterConf.Sources)) |
| for _, dataSource := range ingesterConf.Sources { |
| oneSource, err := getSource(id, dataSource, client, eventBus) |
| if err != nil { |
| return nil, skerr.Wrapf(err, "Error instantiating sources for ingester %q", id) |
| } |
| sources = append(sources, oneSource) |
| sklog.Infof("Source %s created for ingester %s", oneSource.ID(), id) |
| } |
| |
| // instantiate the processor |
| processor, err := processorConstructor(ctx, vcs, ingesterConf, client) |
| if err != nil { |
| return nil, skerr.Wrapf(err, "could not create processor %q", id) |
| } |
| sklog.Infof("Processor constructor for ingester %s created", id) |
| |
| // create the ingester and add it to the result. |
| ingester, err := NewIngester(id, ingesterConf, vcs, sources, processor, ingestionStore, eventBus) |
| if err != nil { |
| return nil, skerr.Wrapf(err, "could not create ingester %q", id) |
| } |
| ret = append(ret, ingester) |
| sklog.Infof("Ingester %s created successfully", id) |
| } |
| |
| return ret, nil |
| } |
| |
| // getSource returns an instance of source that is either getting data from |
| // Google storage or the local filesystem. |
| func getSource(id string, dataSource *sharedconfig.DataSource, client *http.Client, eventBus eventbus.EventBus) (Source, error) { |
| if dataSource.Dir == "" { |
| return nil, fmt.Errorf("Datasource for %s is missing a directory.", id) |
| } |
| |
| if dataSource.Bucket != "" { |
| return NewGoogleStorageSource(id, dataSource.Bucket, dataSource.Dir, client, eventBus) |
| } |
| |
| return nil, skerr.Fmt("Unable to create source. At least a bucket and directory must be supplied") |
| } |
| |
| // validIngestionFile returns true if the given file name matches basic rules. |
| func validIngestionFile(fName string) bool { |
| return targetFileRegExp.Match([]byte(fName)) |
| } |
| |
| // targetFileRegExp must be matched for a file to be considered for ingestion. |
| var targetFileRegExp = regexp.MustCompile(`.*\.json`) |
| |
| // GoogleStorageSource implements the Source interface for Google Storage. |
| type GoogleStorageSource struct { |
| bucket string |
| rootDir string |
| id string |
| storageClient *storage.Client |
| eventBus eventbus.EventBus |
| } |
| |
| // NewGoogleStorageSource returns a new instance of GoogleStorageSource based |
| // on the bucket and directory provided. The id is used to identify the Source |
| // and is generally the same id as the ingester. |
| func NewGoogleStorageSource(baseName, bucket, rootDir string, client *http.Client, eventBus eventbus.EventBus) (Source, error) { |
| storageClient, err := storage.NewClient(context.Background(), option.WithHTTPClient(client)) |
| if err != nil { |
| return nil, skerr.Fmt("Failed to create a Google Storage API client: %s", err) |
| } |
| |
| return &GoogleStorageSource{ |
| bucket: bucket, |
| rootDir: rootDir, |
| id: fmt.Sprintf("%s:gs://%s/%s", baseName, bucket, rootDir), |
| storageClient: storageClient, |
| eventBus: eventBus, |
| }, nil |
| } |
| |
| // See Source interface. |
| func (g *GoogleStorageSource) Poll(startTime, endTime int64) <-chan ResultFileLocation { |
| dirs := fileutil.GetHourlyDirs(g.rootDir, startTime, endTime) |
| ch := make(chan ResultFileLocation, maxConcurrentDirPollers) |
| concurrentPollers := make(chan bool, maxConcurrentDirPollers) |
| |
| go func() { |
| var wg sync.WaitGroup |
| for _, dir := range dirs { |
| concurrentPollers <- true |
| wg.Add(1) |
| |
| go func(dir string) { |
| defer func() { |
| <-concurrentPollers |
| wg.Done() |
| }() |
| err := gcs.AllFilesInDir(g.storageClient, g.bucket, dir, func(item *storage.ObjectAttrs) { |
| if validIngestionFile(item.Name) && (item.Updated.Unix() > startTime) { |
| ch <- newGCSResultFileLocation(item.Bucket, |
| item.Name, |
| item.Updated.Unix(), |
| hex.EncodeToString(item.MD5), |
| g.storageClient) |
| } |
| }) |
| |
| if err != nil { |
| sklog.Errorf("Error occurred while retrieving files from %s/%s: %s", g.bucket, dir, err) |
| } |
| }(dir) |
| } |
| wg.Wait() |
| close(ch) |
| }() |
| |
| return ch |
| } |
| |
| // See Source interface. |
| func (g *GoogleStorageSource) ID() string { |
| return g.id |
| } |
| |
| // SetEventChannel implements the Source interface. |
| func (g *GoogleStorageSource) SetEventChannel(resultCh chan<- ResultFileLocation) error { |
| if g.eventBus != nil { |
| eventType, err := g.eventBus.RegisterStorageEvents(g.bucket, g.rootDir, targetFileRegExp, g.storageClient) |
| if err != nil { |
| return skerr.Fmt("Unable to register storage event: %s", err) |
| } |
| |
| g.eventBus.SubscribeAsync(eventType, func(evData interface{}) { |
| file := evData.(*eventbus.StorageEvent) |
| resultCh <- newGCSResultFileLocation(file.BucketID, file.ObjectID, file.TimeStamp, file.MD5, g.storageClient) |
| }) |
| sklog.Infof("Registered for storage event type: %q", eventType) |
| } |
| return nil |
| } |
| |
| // gsResultFileLocation implements the ResultFileLocation for Google storage. |
| type gsResultFileLocation struct { |
| bucket string |
| name string |
| lastUpdated int64 |
| md5 string |
| storageClient *storage.Client |
| content []byte |
| } |
| |
| func newGCSResultFileLocation(bucketID, objectID string, lastUpdated int64, md5 string, storageClient *storage.Client) ResultFileLocation { |
| return &gsResultFileLocation{ |
| bucket: bucketID, |
| name: objectID, |
| lastUpdated: lastUpdated, |
| md5: md5, |
| storageClient: storageClient, |
| } |
| } |
| |
| // See ResultFileLocation interface. |
| func (g *gsResultFileLocation) Open() (io.ReadCloser, error) { |
| // If we have read this before, then just return a reader. |
| if g.content != nil { |
| return ioutil.NopCloser(bytes.NewBuffer(g.content)), nil |
| } |
| |
| exp := &backoff.ExponentialBackOff{ |
| InitialInterval: time.Second, |
| RandomizationFactor: 0.5, |
| Multiplier: 2, |
| MaxInterval: maxReadTime / 4, |
| MaxElapsedTime: maxReadTime, |
| Clock: backoff.SystemClock, |
| } |
| |
| o := func() error { |
| obj := g.storageClient.Bucket(g.bucket).Object(g.name) |
| reader, err := obj.NewReader(context.Background()) |
| if err != nil { |
| return skerr.Fmt("accessing %s/%s failed: %s", g.bucket, g.name, err) |
| } |
| defer util.Close(reader) |
| |
| // Read the entire file into memory and return a buffer. |
| if g.content, err = ioutil.ReadAll(reader); err != nil { |
| g.content = nil |
| return skerr.Fmt("error reading content of %s/%s: %s", g.bucket, g.name, err) |
| } |
| |
| if oa, err := obj.Attrs(context.Background()); err != nil { |
| g.content = nil |
| return skerr.Fmt("error reading attributes of %s/%s: %s", g.bucket, g.name, err) |
| } else { |
| g.md5 = fmt.Sprintf("%x", oa.MD5) |
| g.lastUpdated = oa.Updated.Unix() |
| } |
| |
| sklog.Debugf("Ingester read %s/%s", g.bucket, g.name) |
| return nil |
| } |
| |
| if g.content == nil { |
| if err := backoff.Retry(o, exp); err != nil { |
| return nil, skerr.Fmt("could not read gcs://%s/%s with retries: %s", g.bucket, g.name, err) |
| } |
| } |
| |
| return ioutil.NopCloser(bytes.NewBuffer(g.content)), nil |
| } |
| |
| // See ResultFileLocation interface. |
| func (g *gsResultFileLocation) Name() string { |
| return fmt.Sprintf("gs://%s/%s", g.bucket, g.name) |
| } |
| |
| // StorageIDs implements the ResultFileLocation interface. |
| func (g *gsResultFileLocation) StorageIDs() (string, string) { |
| return g.bucket, g.name |
| } |
| |
| // See ResultFileLocation interface. |
| func (g *gsResultFileLocation) MD5() string { |
| return g.md5 |
| } |
| |
| // See ResultFileLocation interface. |
| func (g *gsResultFileLocation) TimeStamp() int64 { |
| return g.lastUpdated |
| } |
| |
| // See ResultFileLocation interface. |
| func (g *gsResultFileLocation) Content() []byte { |
| return g.content |
| } |