1051 lines · 27.6 KB · Go
RawDownload
// +build !plan9,!js package cache import (	"bytes"	"context"	"encoding/binary"	"encoding/json"	"fmt"	"io/ioutil"	"os"	"path"	"strconv"	"strings"	"sync"	"time" 	"github.com/pkg/errors"	"github.com/rclone/rclone/fs"	"github.com/rclone/rclone/fs/walk"	bolt "go.etcd.io/bbolt") // Constantsconst (	RootBucket   = "root"	RootTsBucket = "rootTs"	DataTsBucket = "dataTs"	tempBucket   = "pending") // Features flags for this storage typetype Features struct {	PurgeDb    bool          // purge the db before starting	DbWaitTime time.Duration // time to wait for DB to be available} var boltMap = make(map[string]*Persistent)var boltMapMx sync.Mutex // GetPersistent returns a single instance for the specific storefunc GetPersistent(dbPath, chunkPath string, f *Features) (*Persistent, error) {	// write lock to create one	boltMapMx.Lock()	defer boltMapMx.Unlock()	if b, ok := boltMap[dbPath]; ok {		if !b.open {			err := b.connect()			if err != nil {				return nil, err			}		}		return b, nil	} 	bb, err := newPersistent(dbPath, chunkPath, f)	if err != nil {		return nil, err	}	boltMap[dbPath] = bb	return boltMap[dbPath], nil} type chunkInfo struct {	Path   string	Offset int64	Size   int64} type tempUploadInfo struct {	DestPath string	AddedOn  time.Time	Started  bool} // String representation of a tempUploadInfofunc (t *tempUploadInfo) String() string {	return fmt.Sprintf("%v - %v (%v)", t.DestPath, t.Started, t.AddedOn)} // Persistent is a wrapper of persistent storage for a bolt.DB filetype Persistent struct {	dbPath       string	dataPath     string	open         bool	db           *bolt.DB	cleanupMux   sync.Mutex	tempQueueMux sync.Mutex	features     *Features} // newPersistent builds a new wrapper and connects to the bolt.DB filefunc newPersistent(dbPath, chunkPath string, f *Features) (*Persistent, error) {	b := &Persistent{		dbPath:   dbPath,		dataPath: chunkPath,		features: f,	} 	err := b.connect()	if err != nil {		fs.Errorf(dbPath, "Error opening storage cache. Is there another rclone running on the same remote? %v", err)		return nil, err	} 	return b, nil} // String will return a human friendly string for this DB (currently the dbPath)func (b *Persistent) String() string {	return "<Cache DB> " + b.dbPath} // connect creates a connection to the configured file// refreshDb will delete the file before to create an empty DB if it's set to truefunc (b *Persistent) connect() error {	var err error 	err = os.MkdirAll(b.dataPath, os.ModePerm)	if err != nil {		return errors.Wrapf(err, "failed to create a data directory %q", b.dataPath)	}	b.db, err = bolt.Open(b.dbPath, 0644, &bolt.Options{Timeout: b.features.DbWaitTime})	if err != nil {		return errors.Wrapf(err, "failed to open a cache connection to %q", b.dbPath)	}	if b.features.PurgeDb {		b.Purge()	}	_ = b.db.Update(func(tx *bolt.Tx) error {		_, _ = tx.CreateBucketIfNotExists([]byte(RootBucket))		_, _ = tx.CreateBucketIfNotExists([]byte(RootTsBucket))		_, _ = tx.CreateBucketIfNotExists([]byte(DataTsBucket))		_, _ = tx.CreateBucketIfNotExists([]byte(tempBucket)) 		return nil	}) 	b.open = true	return nil} // getBucket prepares and cleans a specific path of the form: /var/tmp and will iterate through each path component// to get to the nested bucket of the final part (in this example: tmp)func (b *Persistent) getBucket(dir string, createIfMissing bool, tx *bolt.Tx) *bolt.Bucket {	cleanPath(dir) 	entries := strings.FieldsFunc(dir, func(c rune) bool {		// cover Windows where rclone still uses '/' as path separator		// this should be safe as '/' is not a valid Windows character		return (os.PathSeparator == c || c == rune('/'))	})	bucket := tx.Bucket([]byte(RootBucket)) 	for _, entry := range entries {		if createIfMissing {			bucket, _ = bucket.CreateBucketIfNotExists([]byte(entry))		} else {			bucket = bucket.Bucket([]byte(entry))		} 		if bucket == nil {			return nil		}	} 	return bucket} // GetDir will retrieve data of a cached directoryfunc (b *Persistent) GetDir(remote string) (*Directory, error) {	cd := &Directory{} 	err := b.db.View(func(tx *bolt.Tx) error {		bucket := b.getBucket(remote, false, tx)		if bucket == nil {			return errors.Errorf("couldn't open bucket (%v)", remote)		} 		data := bucket.Get([]byte("."))		if data != nil {			return json.Unmarshal(data, cd)		} 		return errors.Errorf("%v not found", remote)	}) 	return cd, err} // AddDir will update a CachedDirectory metadata and all its entriesfunc (b *Persistent) AddDir(cachedDir *Directory) error {	return b.AddBatchDir([]*Directory{cachedDir})} // AddBatchDir will update a list of CachedDirectory metadata and all their entriesfunc (b *Persistent) AddBatchDir(cachedDirs []*Directory) error {	if len(cachedDirs) == 0 {		return nil	} 	return b.db.Update(func(tx *bolt.Tx) error {		var bucket *bolt.Bucket		if cachedDirs[0].Dir == "" {			bucket = tx.Bucket([]byte(RootBucket))		} else {			bucket = b.getBucket(cachedDirs[0].Dir, true, tx)		}		if bucket == nil {			return errors.Errorf("couldn't open bucket (%v)", cachedDirs[0].Dir)		} 		for _, cachedDir := range cachedDirs {			var b *bolt.Bucket			var err error			if cachedDir.Name == "" {				b = bucket			} else {				b, err = bucket.CreateBucketIfNotExists([]byte(cachedDir.Name))			}			if err != nil {				return err			} 			encoded, err := json.Marshal(cachedDir)			if err != nil {				return errors.Errorf("couldn't marshal object (%v): %v", cachedDir, err)			}			err = b.Put([]byte("."), encoded)			if err != nil {				return err			}		}		return nil	})} // GetDirEntries will return a CachedDirectory, its list of dir entries and/or an error if it encountered issuesfunc (b *Persistent) GetDirEntries(cachedDir *Directory) (fs.DirEntries, error) {	var dirEntries fs.DirEntries 	err := b.db.View(func(tx *bolt.Tx) error {		bucket := b.getBucket(cachedDir.abs(), false, tx)		if bucket == nil {			return errors.Errorf("couldn't open bucket (%v)", cachedDir.abs())		} 		val := bucket.Get([]byte("."))		if val != nil {			err := json.Unmarshal(val, cachedDir)			if err != nil {				return errors.Errorf("error during unmarshalling obj: %v", err)			}		} else {			return errors.Errorf("missing cached dir: %v", cachedDir)		} 		c := bucket.Cursor()		for k, v := c.First(); k != nil; k, v = c.Next() {			// ignore metadata key: .			if bytes.Equal(k, []byte(".")) {				continue			}			entryPath := path.Join(cachedDir.Remote(), string(k)) 			if v == nil { // directory				// we try to find a cached meta for the dir				currentBucket := c.Bucket().Bucket(k)				if currentBucket == nil {					return errors.Errorf("couldn't open bucket (%v)", string(k))				} 				metaKey := currentBucket.Get([]byte("."))				d := NewDirectory(cachedDir.CacheFs, entryPath)				if metaKey != nil { //if we don't find it, we create an empty dir					err := json.Unmarshal(metaKey, d)					if err != nil { // if even this fails, we fallback to an empty dir						fs.Debugf(string(k), "error during unmarshalling obj: %v", err)					}				} 				dirEntries = append(dirEntries, d)			} else { // object				o := NewObject(cachedDir.CacheFs, entryPath)				err := json.Unmarshal(v, o)				if err != nil {					fs.Debugf(string(k), "error during unmarshalling obj: %v", err)					continue				} 				dirEntries = append(dirEntries, o)			}		} 		return nil	}) 	return dirEntries, err} // RemoveDir will delete a CachedDirectory, all its objects and all the chunks stored for itfunc (b *Persistent) RemoveDir(fp string) error {	var err error	parentDir, dirName := path.Split(fp)	if fp == "" {		err = b.db.Update(func(tx *bolt.Tx) error {			err := tx.DeleteBucket([]byte(RootBucket))			if err != nil {				fs.Debugf(fp, "couldn't delete from cache: %v", err)				return err			}			_, _ = tx.CreateBucketIfNotExists([]byte(RootBucket))			return nil		})	} else {		err = b.db.Update(func(tx *bolt.Tx) error {			bucket := b.getBucket(cleanPath(parentDir), false, tx)			if bucket == nil {				return errors.Errorf("couldn't open bucket (%v)", fp)			}			// delete the cached dir			err := bucket.DeleteBucket([]byte(cleanPath(dirName)))			if err != nil {				fs.Debugf(fp, "couldn't delete from cache: %v", err)			}			return nil		})	} 	// delete chunks on disk	// safe to ignore as the files might not have been open	if err == nil {		_ = os.RemoveAll(path.Join(b.dataPath, fp))		_ = os.MkdirAll(b.dataPath, os.ModePerm)	} 	return err} // ExpireDir will flush a CachedDirectory and all its objects from the objects// chunks will remain as they arefunc (b *Persistent) ExpireDir(cd *Directory) error {	t := time.Now().Add(time.Duration(-cd.CacheFs.opt.InfoAge))	cd.CacheTs = &t 	// expire all parents	return b.db.Update(func(tx *bolt.Tx) error {		// expire all the parents		currentDir := cd.abs()		for { // until we get to the root			bucket := b.getBucket(currentDir, false, tx)			if bucket != nil {				val := bucket.Get([]byte("."))				if val != nil {					cd2 := &Directory{CacheFs: cd.CacheFs}					err := json.Unmarshal(val, cd2)					if err == nil {						fs.Debugf(cd, "cache: expired %v", currentDir)						cd2.CacheTs = &t						enc2, _ := json.Marshal(cd2)						_ = bucket.Put([]byte("."), enc2)					}				}			}			if currentDir == "" {				break			}			currentDir = cleanPath(path.Dir(currentDir))		}		return nil	})} // GetObject will return a CachedObject from its parent directory or an error if it doesn't find itfunc (b *Persistent) GetObject(cachedObject *Object) (err error) {	return b.db.View(func(tx *bolt.Tx) error {		bucket := b.getBucket(cachedObject.Dir, false, tx)		if bucket == nil {			return errors.Errorf("couldn't open parent bucket for %v", cachedObject.Dir)		}		val := bucket.Get([]byte(cachedObject.Name))		if val != nil {			return json.Unmarshal(val, cachedObject)		}		return errors.Errorf("couldn't find object (%v)", cachedObject.Name)	})} // AddObject will create a cached object in its parent directoryfunc (b *Persistent) AddObject(cachedObject *Object) error {	return b.db.Update(func(tx *bolt.Tx) error {		bucket := b.getBucket(cachedObject.Dir, true, tx)		if bucket == nil {			return errors.Errorf("couldn't open parent bucket for %v", cachedObject)		}		// cache Object Info		encoded, err := json.Marshal(cachedObject)		if err != nil {			return errors.Errorf("couldn't marshal object (%v) info: %v", cachedObject, err)		}		err = bucket.Put([]byte(cachedObject.Name), encoded)		if err != nil {			return errors.Errorf("couldn't cache object (%v) info: %v", cachedObject, err)		}		return nil	})} // RemoveObject will delete a single cached object and all the chunks which belong to itfunc (b *Persistent) RemoveObject(fp string) error {	parentDir, objName := path.Split(fp)	return b.db.Update(func(tx *bolt.Tx) error {		bucket := b.getBucket(cleanPath(parentDir), false, tx)		if bucket == nil {			return errors.Errorf("couldn't open parent bucket for %v", cleanPath(parentDir))		}		err := bucket.Delete([]byte(cleanPath(objName)))		if err != nil {			fs.Debugf(fp, "couldn't delete obj from storage: %v", err)		}		// delete chunks on disk		// safe to ignore as the file might not have been open		_ = os.RemoveAll(path.Join(b.dataPath, fp))		return nil	})} // ExpireObject will flush an Object and all its data if desiredfunc (b *Persistent) ExpireObject(co *Object, withData bool) error {	co.CacheTs = time.Now().Add(time.Duration(-co.CacheFs.opt.InfoAge))	err := b.AddObject(co)	if withData {		_ = os.RemoveAll(path.Join(b.dataPath, co.abs()))	}	return err} // HasEntry confirms the existence of a single entry (dir or object)func (b *Persistent) HasEntry(remote string) bool {	dir, name := path.Split(remote)	dir = cleanPath(dir)	name = cleanPath(name) 	err := b.db.View(func(tx *bolt.Tx) error {		bucket := b.getBucket(dir, false, tx)		if bucket == nil {			return errors.Errorf("couldn't open parent bucket for %v", remote)		}		if f := bucket.Bucket([]byte(name)); f != nil {			return nil		}		if f := bucket.Get([]byte(name)); f != nil {			return nil		} 		return errors.Errorf("couldn't find object (%v)", remote)	})	if err == nil {		return true	}	return false} // HasChunk confirms the existence of a single chunk of an objectfunc (b *Persistent) HasChunk(cachedObject *Object, offset int64) bool {	fp := path.Join(b.dataPath, cachedObject.abs(), strconv.FormatInt(offset, 10))	if _, err := os.Stat(fp); !os.IsNotExist(err) {		return true	}	return false} // GetChunk will retrieve a single chunk which belongs to a cached object or an error if it doesn't find itfunc (b *Persistent) GetChunk(cachedObject *Object, offset int64) ([]byte, error) {	var data []byte 	fp := path.Join(b.dataPath, cachedObject.abs(), strconv.FormatInt(offset, 10))	data, err := ioutil.ReadFile(fp)	if err != nil {		return nil, err	} 	return data, err} // AddChunk adds a new chunk of a cached objectfunc (b *Persistent) AddChunk(fp string, data []byte, offset int64) error {	_ = os.MkdirAll(path.Join(b.dataPath, fp), os.ModePerm) 	filePath := path.Join(b.dataPath, fp, strconv.FormatInt(offset, 10))	err := ioutil.WriteFile(filePath, data, os.ModePerm)	if err != nil {		return err	} 	return b.db.Update(func(tx *bolt.Tx) error {		tsBucket := tx.Bucket([]byte(DataTsBucket))		ts := time.Now()		found := false

Lines 1–500 of 1051.