cmd/bpa-operator/vendor/github.com/peterbourgon/diskv/diskv.go

   1 // Diskv (disk-vee) is a simple, persistent, key-value store.
   2 // It stores all data flatly on the filesystem.
   3
   4 package diskv
   5
   6 import (
   7         "bytes"
   8         "errors"
   9         "fmt"
  10         "io"
  11         "io/ioutil"
  12         "os"
  13         "path/filepath"
  14         "strings"
  15         "sync"
  16         "syscall"
  17 )
  18
  19 const (
  20         defaultBasePath             = "diskv"
  21         defaultFilePerm os.FileMode = 0666
  22         defaultPathPerm os.FileMode = 0777
  23 )
  24
  25 var (
  26         defaultTransform   = func(s string) []string { return []string{} }
  27         errCanceled        = errors.New("canceled")
  28         errEmptyKey        = errors.New("empty key")
  29         errBadKey          = errors.New("bad key")
  30         errImportDirectory = errors.New("can't import a directory")
  31 )
  32
  33 // TransformFunction transforms a key into a slice of strings, with each
  34 // element in the slice representing a directory in the file path where the
  35 // key's entry will eventually be stored.
  36 //
  37 // For example, if TransformFunc transforms "abcdef" to ["ab", "cde", "f"],
  38 // the final location of the data file will be <basedir>/ab/cde/f/abcdef
  39 type TransformFunction func(s string) []string
  40
  41 // Options define a set of properties that dictate Diskv behavior.
  42 // All values are optional.
  43 type Options struct {
  44         BasePath     string
  45         Transform    TransformFunction
  46         CacheSizeMax uint64 // bytes
  47         PathPerm     os.FileMode
  48         FilePerm     os.FileMode
  49         // If TempDir is set, it will enable filesystem atomic writes by
  50         // writing temporary files to that location before being moved
  51         // to BasePath.
  52         // Note that TempDir MUST be on the same device/partition as
  53         // BasePath.
  54         TempDir string
  55
  56         Index     Index
  57         IndexLess LessFunction
  58
  59         Compression Compression
  60 }
  61
  62 // Diskv implements the Diskv interface. You shouldn't construct Diskv
  63 // structures directly; instead, use the New constructor.
  64 type Diskv struct {
  65         Options
  66         mu        sync.RWMutex
  67         cache     map[string][]byte
  68         cacheSize uint64
  69 }
  70
  71 // New returns an initialized Diskv structure, ready to use.
  72 // If the path identified by baseDir already contains data,
  73 // it will be accessible, but not yet cached.
  74 func New(o Options) *Diskv {
  75         if o.BasePath == "" {
  76                 o.BasePath = defaultBasePath
  77         }
  78         if o.Transform == nil {
  79                 o.Transform = defaultTransform
  80         }
  81         if o.PathPerm == 0 {
  82                 o.PathPerm = defaultPathPerm
  83         }
  84         if o.FilePerm == 0 {
  85                 o.FilePerm = defaultFilePerm
  86         }
  87
  88         d := &Diskv{
  89                 Options:   o,
  90                 cache:     map[string][]byte{},
  91                 cacheSize: 0,
  92         }
  93
  94         if d.Index != nil && d.IndexLess != nil {
  95                 d.Index.Initialize(d.IndexLess, d.Keys(nil))
  96         }
  97
  98         return d
  99 }
 100
 101 // Write synchronously writes the key-value pair to disk, making it immediately
 102 // available for reads. Write relies on the filesystem to perform an eventual
 103 // sync to physical media. If you need stronger guarantees, see WriteStream.
 104 func (d *Diskv) Write(key string, val []byte) error {
 105         return d.WriteStream(key, bytes.NewBuffer(val), false)
 106 }
 107
 108 // WriteStream writes the data represented by the io.Reader to the disk, under
 109 // the provided key. If sync is true, WriteStream performs an explicit sync on
 110 // the file as soon as it's written.
 111 //
 112 // bytes.Buffer provides io.Reader semantics for basic data types.
 113 func (d *Diskv) WriteStream(key string, r io.Reader, sync bool) error {
 114         if len(key) <= 0 {
 115                 return errEmptyKey
 116         }
 117
 118         d.mu.Lock()
 119         defer d.mu.Unlock()
 120
 121         return d.writeStreamWithLock(key, r, sync)
 122 }
 123
 124 // createKeyFileWithLock either creates the key file directly, or
 125 // creates a temporary file in TempDir if it is set.
 126 func (d *Diskv) createKeyFileWithLock(key string) (*os.File, error) {
 127         if d.TempDir != "" {
 128                 if err := os.MkdirAll(d.TempDir, d.PathPerm); err != nil {
 129                         return nil, fmt.Errorf("temp mkdir: %s", err)
 130                 }
 131                 f, err := ioutil.TempFile(d.TempDir, "")
 132                 if err != nil {
 133                         return nil, fmt.Errorf("temp file: %s", err)
 134                 }
 135
 136                 if err := f.Chmod(d.FilePerm); err != nil {
 137                         f.Close()           // error deliberately ignored
 138                         os.Remove(f.Name()) // error deliberately ignored
 139                         return nil, fmt.Errorf("chmod: %s", err)
 140                 }
 141                 return f, nil
 142         }
 143
 144         mode := os.O_WRONLY | os.O_CREATE | os.O_TRUNC // overwrite if exists
 145         f, err := os.OpenFile(d.completeFilename(key), mode, d.FilePerm)
 146         if err != nil {
 147                 return nil, fmt.Errorf("open file: %s", err)
 148         }
 149         return f, nil
 150 }
 151
 152 // writeStream does no input validation checking.
 153 func (d *Diskv) writeStreamWithLock(key string, r io.Reader, sync bool) error {
 154         if err := d.ensurePathWithLock(key); err != nil {
 155                 return fmt.Errorf("ensure path: %s", err)
 156         }
 157
 158         f, err := d.createKeyFileWithLock(key)
 159         if err != nil {
 160                 return fmt.Errorf("create key file: %s", err)
 161         }
 162
 163         wc := io.WriteCloser(&nopWriteCloser{f})
 164         if d.Compression != nil {
 165                 wc, err = d.Compression.Writer(f)
 166                 if err != nil {
 167                         f.Close()           // error deliberately ignored
 168                         os.Remove(f.Name()) // error deliberately ignored
 169                         return fmt.Errorf("compression writer: %s", err)
 170                 }
 171         }
 172
 173         if _, err := io.Copy(wc, r); err != nil {
 174                 f.Close()           // error deliberately ignored
 175                 os.Remove(f.Name()) // error deliberately ignored
 176                 return fmt.Errorf("i/o copy: %s", err)
 177         }
 178
 179         if err := wc.Close(); err != nil {
 180                 f.Close()           // error deliberately ignored
 181                 os.Remove(f.Name()) // error deliberately ignored
 182                 return fmt.Errorf("compression close: %s", err)
 183         }
 184
 185         if sync {
 186                 if err := f.Sync(); err != nil {
 187                         f.Close()           // error deliberately ignored
 188                         os.Remove(f.Name()) // error deliberately ignored
 189                         return fmt.Errorf("file sync: %s", err)
 190                 }
 191         }
 192
 193         if err := f.Close(); err != nil {
 194                 return fmt.Errorf("file close: %s", err)
 195         }
 196
 197         if f.Name() != d.completeFilename(key) {
 198                 if err := os.Rename(f.Name(), d.completeFilename(key)); err != nil {
 199                         os.Remove(f.Name()) // error deliberately ignored
 200                         return fmt.Errorf("rename: %s", err)
 201                 }
 202         }
 203
 204         if d.Index != nil {
 205                 d.Index.Insert(key)
 206         }
 207
 208         d.bustCacheWithLock(key) // cache only on read
 209
 210         return nil
 211 }
 212
 213 // Import imports the source file into diskv under the destination key. If the
 214 // destination key already exists, it's overwritten. If move is true, the
 215 // source file is removed after a successful import.
 216 func (d *Diskv) Import(srcFilename, dstKey string, move bool) (err error) {
 217         if dstKey == "" {
 218                 return errEmptyKey
 219         }
 220
 221         if fi, err := os.Stat(srcFilename); err != nil {
 222                 return err
 223         } else if fi.IsDir() {
 224                 return errImportDirectory
 225         }
 226
 227         d.mu.Lock()
 228         defer d.mu.Unlock()
 229
 230         if err := d.ensurePathWithLock(dstKey); err != nil {
 231                 return fmt.Errorf("ensure path: %s", err)
 232         }
 233
 234         if move {
 235                 if err := syscall.Rename(srcFilename, d.completeFilename(dstKey)); err == nil {
 236                         d.bustCacheWithLock(dstKey)
 237                         return nil
 238                 } else if err != syscall.EXDEV {
 239                         // If it failed due to being on a different device, fall back to copying
 240                         return err
 241                 }
 242         }
 243
 244         f, err := os.Open(srcFilename)
 245         if err != nil {
 246                 return err
 247         }
 248         defer f.Close()
 249         err = d.writeStreamWithLock(dstKey, f, false)
 250         if err == nil && move {
 251                 err = os.Remove(srcFilename)
 252         }
 253         return err
 254 }
 255
 256 // Read reads the key and returns the value.
 257 // If the key is available in the cache, Read won't touch the disk.
 258 // If the key is not in the cache, Read will have the side-effect of
 259 // lazily caching the value.
 260 func (d *Diskv) Read(key string) ([]byte, error) {
 261         rc, err := d.ReadStream(key, false)
 262         if err != nil {
 263                 return []byte{}, err
 264         }
 265         defer rc.Close()
 266         return ioutil.ReadAll(rc)
 267 }
 268
 269 // ReadStream reads the key and returns the value (data) as an io.ReadCloser.
 270 // If the value is cached from a previous read, and direct is false,
 271 // ReadStream will use the cached value. Otherwise, it will return a handle to
 272 // the file on disk, and cache the data on read.
 273 //
 274 // If direct is true, ReadStream will lazily delete any cached value for the
 275 // key, and return a direct handle to the file on disk.
 276 //
 277 // If compression is enabled, ReadStream taps into the io.Reader stream prior
 278 // to decompression, and caches the compressed data.
 279 func (d *Diskv) ReadStream(key string, direct bool) (io.ReadCloser, error) {
 280         d.mu.RLock()
 281         defer d.mu.RUnlock()
 282
 283         if val, ok := d.cache[key]; ok {
 284                 if !direct {
 285                         buf := bytes.NewBuffer(val)
 286                         if d.Compression != nil {
 287                                 return d.Compression.Reader(buf)
 288                         }
 289                         return ioutil.NopCloser(buf), nil
 290                 }
 291
 292                 go func() {
 293                         d.mu.Lock()
 294                         defer d.mu.Unlock()
 295                         d.uncacheWithLock(key, uint64(len(val)))
 296                 }()
 297         }
 298
 299         return d.readWithRLock(key)
 300 }
 301
 302 // read ignores the cache, and returns an io.ReadCloser representing the
 303 // decompressed data for the given key, streamed from the disk. Clients should
 304 // acquire a read lock on the Diskv and check the cache themselves before
 305 // calling read.
 306 func (d *Diskv) readWithRLock(key string) (io.ReadCloser, error) {
 307         filename := d.completeFilename(key)
 308
 309         fi, err := os.Stat(filename)
 310         if err != nil {
 311                 return nil, err
 312         }
 313         if fi.IsDir() {
 314                 return nil, os.ErrNotExist
 315         }
 316
 317         f, err := os.Open(filename)
 318         if err != nil {
 319                 return nil, err
 320         }
 321
 322         var r io.Reader
 323         if d.CacheSizeMax > 0 {
 324                 r = newSiphon(f, d, key)
 325         } else {
 326                 r = &closingReader{f}
 327         }
 328
 329         var rc = io.ReadCloser(ioutil.NopCloser(r))
 330         if d.Compression != nil {
 331                 rc, err = d.Compression.Reader(r)
 332                 if err != nil {
 333                         return nil, err
 334                 }
 335         }
 336
 337         return rc, nil
 338 }
 339
 340 // closingReader provides a Reader that automatically closes the
 341 // embedded ReadCloser when it reaches EOF
 342 type closingReader struct {
 343         rc io.ReadCloser
 344 }
 345
 346 func (cr closingReader) Read(p []byte) (int, error) {
 347         n, err := cr.rc.Read(p)
 348         if err == io.EOF {
 349                 if closeErr := cr.rc.Close(); closeErr != nil {
 350                         return n, closeErr // close must succeed for Read to succeed
 351                 }
 352         }
 353         return n, err
 354 }
 355
 356 // siphon is like a TeeReader: it copies all data read through it to an
 357 // internal buffer, and moves that buffer to the cache at EOF.
 358 type siphon struct {
 359         f   *os.File
 360         d   *Diskv
 361         key string
 362         buf *bytes.Buffer
 363 }
 364
 365 // newSiphon constructs a siphoning reader that represents the passed file.
 366 // When a successful series of reads ends in an EOF, the siphon will write
 367 // the buffered data to Diskv's cache under the given key.
 368 func newSiphon(f *os.File, d *Diskv, key string) io.Reader {
 369         return &siphon{
 370                 f:   f,
 371                 d:   d,
 372                 key: key,
 373                 buf: &bytes.Buffer{},
 374         }
 375 }
 376
 377 // Read implements the io.Reader interface for siphon.
 378 func (s *siphon) Read(p []byte) (int, error) {
 379         n, err := s.f.Read(p)
 380
 381         if err == nil {
 382                 return s.buf.Write(p[0:n]) // Write must succeed for Read to succeed
 383         }
 384
 385         if err == io.EOF {
 386                 s.d.cacheWithoutLock(s.key, s.buf.Bytes()) // cache may fail
 387                 if closeErr := s.f.Close(); closeErr != nil {
 388                         return n, closeErr // close must succeed for Read to succeed
 389                 }
 390                 return n, err
 391         }
 392
 393         return n, err
 394 }
 395
 396 // Erase synchronously erases the given key from the disk and the cache.
 397 func (d *Diskv) Erase(key string) error {
 398         d.mu.Lock()
 399         defer d.mu.Unlock()
 400
 401         d.bustCacheWithLock(key)
 402
 403         // erase from index
 404         if d.Index != nil {
 405                 d.Index.Delete(key)
 406         }
 407
 408         // erase from disk
 409         filename := d.completeFilename(key)
 410         if s, err := os.Stat(filename); err == nil {
 411                 if s.IsDir() {
 412                         return errBadKey
 413                 }
 414                 if err = os.Remove(filename); err != nil {
 415                         return err
 416                 }
 417         } else {
 418                 // Return err as-is so caller can do os.IsNotExist(err).
 419                 return err
 420         }
 421
 422         // clean up and return
 423         d.pruneDirsWithLock(key)
 424         return nil
 425 }
 426
 427 // EraseAll will delete all of the data from the store, both in the cache and on
 428 // the disk. Note that EraseAll doesn't distinguish diskv-related data from non-
 429 // diskv-related data. Care should be taken to always specify a diskv base
 430 // directory that is exclusively for diskv data.
 431 func (d *Diskv) EraseAll() error {
 432         d.mu.Lock()
 433         defer d.mu.Unlock()
 434         d.cache = make(map[string][]byte)
 435         d.cacheSize = 0
 436         if d.TempDir != "" {
 437                 os.RemoveAll(d.TempDir) // errors ignored
 438         }
 439         return os.RemoveAll(d.BasePath)
 440 }
 441
 442 // Has returns true if the given key exists.
 443 func (d *Diskv) Has(key string) bool {
 444         d.mu.Lock()
 445         defer d.mu.Unlock()
 446
 447         if _, ok := d.cache[key]; ok {
 448                 return true
 449         }
 450
 451         filename := d.completeFilename(key)
 452         s, err := os.Stat(filename)
 453         if err != nil {
 454                 return false
 455         }
 456         if s.IsDir() {
 457                 return false
 458         }
 459
 460         return true
 461 }
 462
 463 // Keys returns a channel that will yield every key accessible by the store,
 464 // in undefined order. If a cancel channel is provided, closing it will
 465 // terminate and close the keys channel.
 466 func (d *Diskv) Keys(cancel <-chan struct{}) <-chan string {
 467         return d.KeysPrefix("", cancel)
 468 }
 469
 470 // KeysPrefix returns a channel that will yield every key accessible by the
 471 // store with the given prefix, in undefined order. If a cancel channel is
 472 // provided, closing it will terminate and close the keys channel. If the
 473 // provided prefix is the empty string, all keys will be yielded.
 474 func (d *Diskv) KeysPrefix(prefix string, cancel <-chan struct{}) <-chan string {
 475         var prepath string
 476         if prefix == "" {
 477                 prepath = d.BasePath
 478         } else {
 479                 prepath = d.pathFor(prefix)
 480         }
 481         c := make(chan string)
 482         go func() {
 483                 filepath.Walk(prepath, walker(c, prefix, cancel))
 484                 close(c)
 485         }()
 486         return c
 487 }
 488
 489 // walker returns a function which satisfies the filepath.WalkFunc interface.
 490 // It sends every non-directory file entry down the channel c.
 491 func walker(c chan<- string, prefix string, cancel <-chan struct{}) filepath.WalkFunc {
 492         return func(path string, info os.FileInfo, err error) error {
 493                 if err != nil {
 494                         return err
 495                 }
 496
 497                 if info.IsDir() || !strings.HasPrefix(info.Name(), prefix) {
 498                         return nil // "pass"
 499                 }
 500
 501                 select {
 502                 case c <- info.Name():
 503                 case <-cancel:
 504                         return errCanceled
 505                 }
 506
 507                 return nil
 508         }
 509 }
 510
 511 // pathFor returns the absolute path for location on the filesystem where the
 512 // data for the given key will be stored.
 513 func (d *Diskv) pathFor(key string) string {
 514         return filepath.Join(d.BasePath, filepath.Join(d.Transform(key)...))
 515 }
 516
 517 // ensurePathWithLock is a helper function that generates all necessary
 518 // directories on the filesystem for the given key.
 519 func (d *Diskv) ensurePathWithLock(key string) error {
 520         return os.MkdirAll(d.pathFor(key), d.PathPerm)
 521 }
 522
 523 // completeFilename returns the absolute path to the file for the given key.
 524 func (d *Diskv) completeFilename(key string) string {
 525         return filepath.Join(d.pathFor(key), key)
 526 }
 527
 528 // cacheWithLock attempts to cache the given key-value pair in the store's
 529 // cache. It can fail if the value is larger than the cache's maximum size.
 530 func (d *Diskv) cacheWithLock(key string, val []byte) error {
 531         valueSize := uint64(len(val))
 532         if err := d.ensureCacheSpaceWithLock(valueSize); err != nil {
 533                 return fmt.Errorf("%s; not caching", err)
 534         }
 535
 536         // be very strict about memory guarantees
 537         if (d.cacheSize + valueSize) > d.CacheSizeMax {
 538                 panic(fmt.Sprintf("failed to make room for value (%d/%d)", valueSize, d.CacheSizeMax))
 539         }
 540
 541         d.cache[key] = val
 542         d.cacheSize += valueSize
 543         return nil
 544 }
 545
 546 // cacheWithoutLock acquires the store's (write) mutex and calls cacheWithLock.
 547 func (d *Diskv) cacheWithoutLock(key string, val []byte) error {
 548         d.mu.Lock()
 549         defer d.mu.Unlock()
 550         return d.cacheWithLock(key, val)
 551 }
 552
 553 func (d *Diskv) bustCacheWithLock(key string) {
 554         if val, ok := d.cache[key]; ok {
 555                 d.uncacheWithLock(key, uint64(len(val)))
 556         }
 557 }
 558
 559 func (d *Diskv) uncacheWithLock(key string, sz uint64) {
 560         d.cacheSize -= sz
 561         delete(d.cache, key)
 562 }
 563
 564 // pruneDirsWithLock deletes empty directories in the path walk leading to the
 565 // key k. Typically this function is called after an Erase is made.
 566 func (d *Diskv) pruneDirsWithLock(key string) error {
 567         pathlist := d.Transform(key)
 568         for i := range pathlist {
 569                 dir := filepath.Join(d.BasePath, filepath.Join(pathlist[:len(pathlist)-i]...))
 570
 571                 // thanks to Steven Blenkinsop for this snippet
 572                 switch fi, err := os.Stat(dir); true {
 573                 case err != nil:
 574                         return err
 575                 case !fi.IsDir():
 576                         panic(fmt.Sprintf("corrupt dirstate at %s", dir))
 577                 }
 578
 579                 nlinks, err := filepath.Glob(filepath.Join(dir, "*"))
 580                 if err != nil {
 581                         return err
 582                 } else if len(nlinks) > 0 {
 583                         return nil // has subdirs -- do not prune
 584                 }
 585                 if err = os.Remove(dir); err != nil {
 586                         return err
 587                 }
 588         }
 589
 590         return nil
 591 }
 592
 593 // ensureCacheSpaceWithLock deletes entries from the cache in arbitrary order
 594 // until the cache has at least valueSize bytes available.
 595 func (d *Diskv) ensureCacheSpaceWithLock(valueSize uint64) error {
 596         if valueSize > d.CacheSizeMax {
 597                 return fmt.Errorf("value size (%d bytes) too large for cache (%d bytes)", valueSize, d.CacheSizeMax)
 598         }
 599
 600         safe := func() bool { return (d.cacheSize + valueSize) <= d.CacheSizeMax }
 601
 602         for key, val := range d.cache {
 603                 if safe() {
 604                         break
 605                 }
 606
 607                 d.uncacheWithLock(key, uint64(len(val)))
 608         }
 609
 610         if !safe() {
 611                 panic(fmt.Sprintf("%d bytes still won't fit in the cache! (max %d bytes)", valueSize, d.CacheSizeMax))
 612         }
 613
 614         return nil
 615 }
 616
 617 // nopWriteCloser wraps an io.Writer and provides a no-op Close method to
 618 // satisfy the io.WriteCloser interface.
 619 type nopWriteCloser struct {
 620         io.Writer
 621 }
 622
 623 func (wc *nopWriteCloser) Write(p []byte) (int, error) { return wc.Writer.Write(p) }
 624 func (wc *nopWriteCloser) Close() error                { return nil }