// Code generated by github.com/vektah/dataloaden, DO NOT EDIT. package loaders import ( "sync" "time" "github.com/stashapp/stash/pkg/file" ) // FileLoaderConfig captures the config to create a new FileLoader type FileLoaderConfig struct { // Fetch is a method that provides the data for the loader Fetch func(keys []file.ID) ([]file.File, []error) // Wait is how long wait before sending a batch Wait time.Duration // MaxBatch will limit the maximum number of keys to send in one batch, 0 = not limit MaxBatch int } // NewFileLoader creates a new FileLoader given a fetch, wait, and maxBatch func NewFileLoader(config FileLoaderConfig) *FileLoader { return &FileLoader{ fetch: config.Fetch, wait: config.Wait, maxBatch: config.MaxBatch, } } // FileLoader batches and caches requests type FileLoader struct { // this method provides the data for the loader fetch func(keys []file.ID) ([]file.File, []error) // how long to done before sending a batch wait time.Duration // this will limit the maximum number of keys to send in one batch, 0 = no limit maxBatch int // INTERNAL // lazily created cache cache map[file.ID]file.File // the current batch. keys will continue to be collected until timeout is hit, // then everything will be sent to the fetch method and out to the listeners batch *fileLoaderBatch // mutex to prevent races mu sync.Mutex } type fileLoaderBatch struct { keys []file.ID data []file.File error []error closing bool done chan struct{} } // Load a File by key, batching and caching will be applied automatically func (l *FileLoader) Load(key file.ID) (file.File, error) { return l.LoadThunk(key)() } // LoadThunk returns a function that when called will block waiting for a File. // This method should be used if you want one goroutine to make requests to many // different data loaders without blocking until the thunk is called. func (l *FileLoader) LoadThunk(key file.ID) func() (file.File, error) { l.mu.Lock() if it, ok := l.cache[key]; ok { l.mu.Unlock() return func() (file.File, error) { return it, nil } } if l.batch == nil { l.batch = &fileLoaderBatch{done: make(chan struct{})} } batch := l.batch pos := batch.keyIndex(l, key) l.mu.Unlock() return func() (file.File, error) { <-batch.done var data file.File if pos < len(batch.data) { data = batch.data[pos] } var err error // its convenient to be able to return a single error for everything if len(batch.error) == 1 { err = batch.error[0] } else if batch.error != nil { err = batch.error[pos] } if err == nil { l.mu.Lock() l.unsafeSet(key, data) l.mu.Unlock() } return data, err } } // LoadAll fetches many keys at once. It will be broken into appropriate sized // sub batches depending on how the loader is configured func (l *FileLoader) LoadAll(keys []file.ID) ([]file.File, []error) { results := make([]func() (file.File, error), len(keys)) for i, key := range keys { results[i] = l.LoadThunk(key) } files := make([]file.File, len(keys)) errors := make([]error, len(keys)) for i, thunk := range results { files[i], errors[i] = thunk() } return files, errors } // LoadAllThunk returns a function that when called will block waiting for a Files. // This method should be used if you want one goroutine to make requests to many // different data loaders without blocking until the thunk is called. func (l *FileLoader) LoadAllThunk(keys []file.ID) func() ([]file.File, []error) { results := make([]func() (file.File, error), len(keys)) for i, key := range keys { results[i] = l.LoadThunk(key) } return func() ([]file.File, []error) { files := make([]file.File, len(keys)) errors := make([]error, len(keys)) for i, thunk := range results { files[i], errors[i] = thunk() } return files, errors } } // Prime the cache with the provided key and value. If the key already exists, no change is made // and false is returned. // (To forcefully prime the cache, clear the key first with loader.clear(key).prime(key, value).) func (l *FileLoader) Prime(key file.ID, value file.File) bool { l.mu.Lock() var found bool if _, found = l.cache[key]; !found { l.unsafeSet(key, value) } l.mu.Unlock() return !found } // Clear the value at key from the cache, if it exists func (l *FileLoader) Clear(key file.ID) { l.mu.Lock() delete(l.cache, key) l.mu.Unlock() } func (l *FileLoader) unsafeSet(key file.ID, value file.File) { if l.cache == nil { l.cache = map[file.ID]file.File{} } l.cache[key] = value } // keyIndex will return the location of the key in the batch, if its not found // it will add the key to the batch func (b *fileLoaderBatch) keyIndex(l *FileLoader, key file.ID) int { for i, existingKey := range b.keys { if key == existingKey { return i } } pos := len(b.keys) b.keys = append(b.keys, key) if pos == 0 { go b.startTimer(l) } if l.maxBatch != 0 && pos >= l.maxBatch-1 { if !b.closing { b.closing = true l.batch = nil go b.end(l) } } return pos } func (b *fileLoaderBatch) startTimer(l *FileLoader) { time.Sleep(l.wait) l.mu.Lock() // we must have hit a batch limit and are already finalizing this batch if b.closing { l.mu.Unlock() return } l.batch = nil l.mu.Unlock() b.end(l) } func (b *fileLoaderBatch) end(l *FileLoader) { b.data, b.error = l.fetch(b.keys) close(b.done) }