github.com/uber/kraken@v0.1.4/lib/store/simple_store.go (about)

     1  // Copyright (c) 2016-2019 Uber Technologies, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  package store
    15  
    16  import (
    17  	"fmt"
    18  	"io"
    19  	"os"
    20  
    21  	"github.com/uber/kraken/lib/store/base"
    22  	"github.com/andres-erbsen/clock"
    23  	"github.com/docker/distribution/uuid"
    24  	"github.com/uber-go/tally"
    25  )
    26  
    27  // SimpleStore allows uploading / caching raw files of any format.
    28  type SimpleStore struct {
    29  	*uploadStore
    30  	*cacheStore
    31  	cleanup *cleanupManager
    32  }
    33  
    34  // NewSimpleStore creates a new SimpleStore.
    35  func NewSimpleStore(config SimpleStoreConfig, stats tally.Scope) (*SimpleStore, error) {
    36  	stats = stats.Tagged(map[string]string{
    37  		"module": "simplestore",
    38  	})
    39  
    40  	uploadStore, err := newUploadStore(config.UploadDir)
    41  	if err != nil {
    42  		return nil, fmt.Errorf("new upload store: %s", err)
    43  	}
    44  
    45  	cacheBackend := base.NewLocalFileStore(clock.New())
    46  	cacheStore, err := newCacheStore(config.CacheDir, cacheBackend)
    47  	if err != nil {
    48  		return nil, fmt.Errorf("new cache store: %s", err)
    49  	}
    50  
    51  	cleanup, err := newCleanupManager(clock.New(), stats)
    52  	if err != nil {
    53  		return nil, fmt.Errorf("new cleanup manager: %s", err)
    54  	}
    55  	cleanup.addJob("upload", config.UploadCleanup, uploadStore.newFileOp())
    56  	cleanup.addJob("cache", config.CacheCleanup, cacheStore.newFileOp())
    57  
    58  	return &SimpleStore{uploadStore, cacheStore, cleanup}, nil
    59  }
    60  
    61  // Close terminates goroutines started by s.
    62  func (s *SimpleStore) Close() {
    63  	s.cleanup.stop()
    64  }
    65  
    66  // MoveUploadFileToCache commits uploadName as cacheName.
    67  func (s *SimpleStore) MoveUploadFileToCache(uploadName, cacheName string) error {
    68  	uploadPath, err := s.uploadStore.newFileOp().GetFilePath(uploadName)
    69  	if err != nil {
    70  		return err
    71  	}
    72  	defer s.DeleteUploadFile(uploadName)
    73  	return s.cacheStore.newFileOp().MoveFileFrom(cacheName, s.cacheStore.state, uploadPath)
    74  }
    75  
    76  // CreateCacheFile initializes a cache file for name from r.
    77  func (s *SimpleStore) CreateCacheFile(name string, r io.Reader) error {
    78  	tmp := fmt.Sprintf("%s.%s", name, uuid.Generate().String())
    79  	if err := s.CreateUploadFile(tmp, 0); err != nil {
    80  		return fmt.Errorf("create upload file: %s", err)
    81  	}
    82  	defer s.DeleteUploadFile(tmp)
    83  
    84  	w, err := s.GetUploadFileReadWriter(tmp)
    85  	if err != nil {
    86  		return fmt.Errorf("get upload writer: %s", err)
    87  	}
    88  	defer w.Close()
    89  
    90  	if _, err := io.Copy(w, r); err != nil {
    91  		return fmt.Errorf("copy: %s", err)
    92  	}
    93  
    94  	if err := s.MoveUploadFileToCache(tmp, name); err != nil && !os.IsExist(err) {
    95  		return fmt.Errorf("move upload file to cache: %s", err)
    96  	}
    97  	return nil
    98  }