mount.go 8.28 KB
Newer Older
1 2
// Package mount provides a Datastore that has other Datastores
// mounted at various key prefixes and is threadsafe
Łukasz Magiera's avatar
Łukasz Magiera committed
3
package mount
4 5

import (
Steven Allen's avatar
Steven Allen committed
6
	"container/heap"
7
	"errors"
8
	"sort"
9 10 11
	"strings"
	"sync"

Jeromy's avatar
Jeromy committed
12 13
	ds "github.com/ipfs/go-datastore"
	"github.com/ipfs/go-datastore/query"
14
	xerrors "golang.org/x/xerrors"
15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31
)

var (
	ErrNoMount = errors.New("no datastore mounted for this key")
)

type Mount struct {
	Prefix    ds.Key
	Datastore ds.Datastore
}

func New(mounts []Mount) *Datastore {
	// make a copy so we're sure it doesn't mutate
	m := make([]Mount, len(mounts))
	for i, v := range mounts {
		m[i] = v
	}
32
	sort.Slice(m, func(i, j int) bool { return m[i].Prefix.String() > m[j].Prefix.String() })
33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52
	return &Datastore{mounts: m}
}

type Datastore struct {
	mounts []Mount
}

var _ ds.Datastore = (*Datastore)(nil)

func (d *Datastore) lookup(key ds.Key) (ds.Datastore, ds.Key, ds.Key) {
	for _, m := range d.mounts {
		if m.Prefix.Equal(key) || m.Prefix.IsAncestorOf(key) {
			s := strings.TrimPrefix(key.String(), m.Prefix.String())
			k := ds.NewKey(s)
			return m.Datastore, m.Prefix, k
		}
	}
	return nil, ds.NewKey("/"), key
}

Steven Allen's avatar
Steven Allen committed
53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82
type queryResults struct {
	mount   ds.Key
	results query.Results
	next    query.Result
}

func (qr *queryResults) advance() bool {
	if qr.results == nil {
		return false
	}

	qr.next = query.Result{}
	r, more := qr.results.NextSync()
	if !more {
		err := qr.results.Close()
		qr.results = nil
		if err != nil {
			// One more result, the error.
			qr.next = query.Result{Error: err}
			return true
		}
		return false
	}

	r.Key = qr.mount.Child(ds.RawKey(r.Key)).String()
	qr.next = r
	return true
}

type querySet struct {
Steven Allen's avatar
Steven Allen committed
83
	query query.Query
Steven Allen's avatar
Steven Allen committed
84 85 86 87 88 89 90 91
	heads []*queryResults
}

func (h *querySet) Len() int {
	return len(h.heads)
}

func (h *querySet) Less(i, j int) bool {
Steven Allen's avatar
Steven Allen committed
92
	return query.Less(h.query.Orders, h.heads[i].next.Entry, h.heads[j].next.Entry)
Steven Allen's avatar
Steven Allen committed
93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139
}

func (h *querySet) Swap(i, j int) {
	h.heads[i], h.heads[j] = h.heads[j], h.heads[i]
}

func (h *querySet) Push(x interface{}) {
	h.heads = append(h.heads, x.(*queryResults))
}

func (h *querySet) Pop() interface{} {
	i := len(h.heads) - 1
	last := h.heads[i]
	h.heads[i] = nil
	h.heads = h.heads[:i]
	return last
}

func (h *querySet) close() error {
	var errs []error
	for _, qr := range h.heads {
		err := qr.results.Close()
		if err != nil {
			errs = append(errs, err)
		}
	}
	h.heads = nil
	if len(errs) > 0 {
		return errs[0]
	}
	return nil
}

func (h *querySet) addResults(mount ds.Key, results query.Results) {
	r := &queryResults{
		results: results,
		mount:   mount,
	}
	if r.advance() {
		heap.Push(h, r)
	}
}

func (h *querySet) next() (query.Result, bool) {
	if len(h.heads) == 0 {
		return query.Result{}, false
	}
Steven Allen's avatar
Steven Allen committed
140 141
	head := h.heads[0]
	next := head.next
142

Michael Avila's avatar
Michael Avila committed
143
	if head.advance() {
Steven Allen's avatar
Steven Allen committed
144
		heap.Fix(h, 0)
Michael Avila's avatar
Michael Avila committed
145 146
	} else {
		heap.Remove(h, 0)
Steven Allen's avatar
Steven Allen committed
147
	}
Michael Avila's avatar
Michael Avila committed
148

Steven Allen's avatar
Steven Allen committed
149 150 151
	return next, true
}

152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183
// lookupAll returns all mounts that might contain keys that are descendant of <key>
//
// Matching: /ao/e
//
// /          B /ao/e
// /a/        not matching
// /ao/       B /e
// /ao/e/     A /
// /ao/e/uh/  A /
// /aoe/      not matching
func (d *Datastore) lookupAll(key ds.Key) (dst []ds.Datastore, mountpoint, rest []ds.Key) {
	for _, m := range d.mounts {
		p := m.Prefix.String()
		if len(p) > 1 {
			p = p + "/"
		}

		if strings.HasPrefix(p, key.String()) {
			dst = append(dst, m.Datastore)
			mountpoint = append(mountpoint, m.Prefix)
			rest = append(rest, ds.NewKey("/"))
		} else if strings.HasPrefix(key.String(), p) {
			r := strings.TrimPrefix(key.String(), m.Prefix.String())

			dst = append(dst, m.Datastore)
			mountpoint = append(mountpoint, m.Prefix)
			rest = append(rest, ds.NewKey(r))
		}
	}
	return dst, mountpoint, rest
}

184
func (d *Datastore) Put(key ds.Key, value []byte) error {
185 186 187 188 189 190 191
	cds, _, k := d.lookup(key)
	if cds == nil {
		return ErrNoMount
	}
	return cds.Put(k, value)
}

192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209
// Sync implements Datastore.Sync
func (d *Datastore) Sync(prefix ds.Key) error {
	// Sync all mount points below the prefix
	// Sync the mount point right at (or above) the prefix
	dstores, mountPts, rest := d.lookupAll(prefix)
	for i, suffix := range rest {
		if err := dstores[i].Sync(suffix); err != nil {
			return err
		}

		if mountPts[i].Equal(prefix) || suffix.String() != "/" {
			return nil
		}
	}

	return nil
}

210
func (d *Datastore) Get(key ds.Key) (value []byte, err error) {
211 212 213 214 215 216 217 218 219 220 221 222 223 224 225
	cds, _, k := d.lookup(key)
	if cds == nil {
		return nil, ds.ErrNotFound
	}
	return cds.Get(k)
}

func (d *Datastore) Has(key ds.Key) (exists bool, err error) {
	cds, _, k := d.lookup(key)
	if cds == nil {
		return false, nil
	}
	return cds.Has(k)
}

Steven Allen's avatar
Steven Allen committed
226 227 228 229 230 231 232 233
func (d *Datastore) GetSize(key ds.Key) (size int, err error) {
	cds, _, k := d.lookup(key)
	if cds == nil {
		return -1, ds.ErrNotFound
	}
	return cds.GetSize(k)
}

234 235 236
func (d *Datastore) Delete(key ds.Key) error {
	cds, _, k := d.lookup(key)
	if cds == nil {
Steven Allen's avatar
Steven Allen committed
237
		return nil
238 239 240 241
	}
	return cds.Delete(k)
}

Michael Avila's avatar
Michael Avila committed
242 243
func (d *Datastore) Query(master query.Query) (query.Results, error) {
	childQuery := query.Query{
Michael Avila's avatar
Michael Avila committed
244 245 246
		Prefix:            master.Prefix,
		Orders:            master.Orders,
		KeysOnly:          master.KeysOnly,
Michael Avila's avatar
Michael Avila committed
247
		ReturnExpirations: master.ReturnExpirations,
248
		ReturnsSizes:      master.ReturnsSizes,
Michael Avila's avatar
Michael Avila committed
249
	}
250

Michael Avila's avatar
Michael Avila committed
251 252
	prefix := ds.NewKey(childQuery.Prefix)
	dses, mounts, rests := d.lookupAll(prefix)
253

Steven Allen's avatar
Steven Allen committed
254
	queries := &querySet{
Michael Avila's avatar
Michael Avila committed
255
		query: childQuery,
Steven Allen's avatar
Steven Allen committed
256 257 258 259 260 261 262 263
		heads: make([]*queryResults, 0, len(dses)),
	}

	for i := range dses {
		mount := mounts[i]
		dstore := dses[i]
		rest := rests[i]

Michael Avila's avatar
Michael Avila committed
264
		qi := childQuery
Steven Allen's avatar
Steven Allen committed
265 266
		qi.Prefix = rest.String()
		results, err := dstore.Query(qi)
Steven Allen's avatar
Steven Allen committed
267

Steven Allen's avatar
Steven Allen committed
268 269 270 271 272 273
		if err != nil {
			_ = queries.close()
			return nil, err
		}
		queries.addResults(mount, results)
	}
274

275
	qr := query.ResultsFromIterator(master, query.Iterator{
Steven Allen's avatar
Steven Allen committed
276 277
		Next:  queries.next,
		Close: queries.close,
Michael Avila's avatar
Michael Avila committed
278 279
	})

Michael Avila's avatar
Michael Avila committed
280 281
	if len(master.Filters) > 0 {
		for _, f := range master.Filters {
282 283 284 285
			qr = query.NaiveFilter(qr, f)
		}
	}

Michael Avila's avatar
Michael Avila committed
286 287
	if master.Offset > 0 {
		qr = query.NaiveOffset(qr, master.Offset)
288 289
	}

290 291
	if master.Limit > 0 {
		qr = query.NaiveLimit(qr, master.Limit)
292 293 294
	}

	return qr, nil
295 296 297 298
}

func (d *Datastore) Close() error {
	for _, d := range d.mounts {
Steven Allen's avatar
Steven Allen committed
299 300 301
		err := d.Datastore.Close()
		if err != nil {
			return err
302 303 304 305 306
		}
	}
	return nil
}

307 308 309 310 311 312 313 314 315 316 317 318 319 320
// DiskUsage returns the sum of DiskUsages for the mounted datastores.
// Non PersistentDatastores will not be accounted.
func (d *Datastore) DiskUsage() (uint64, error) {
	var duTotal uint64 = 0
	for _, d := range d.mounts {
		du, err := ds.DiskUsage(d.Datastore)
		duTotal += du
		if err != nil {
			return duTotal, err
		}
	}
	return duTotal, nil
}

321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355
type mountBatch struct {
	mounts map[string]ds.Batch
	lk     sync.Mutex

	d *Datastore
}

func (d *Datastore) Batch() (ds.Batch, error) {
	return &mountBatch{
		mounts: make(map[string]ds.Batch),
		d:      d,
	}, nil
}

func (mt *mountBatch) lookupBatch(key ds.Key) (ds.Batch, ds.Key, error) {
	mt.lk.Lock()
	defer mt.lk.Unlock()

	child, loc, rest := mt.d.lookup(key)
	t, ok := mt.mounts[loc.String()]
	if !ok {
		bds, ok := child.(ds.Batching)
		if !ok {
			return nil, ds.NewKey(""), ds.ErrBatchUnsupported
		}
		var err error
		t, err = bds.Batch()
		if err != nil {
			return nil, ds.NewKey(""), err
		}
		mt.mounts[loc.String()] = t
	}
	return t, rest, nil
}

356
func (mt *mountBatch) Put(key ds.Key, val []byte) error {
357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374
	t, rest, err := mt.lookupBatch(key)
	if err != nil {
		return err
	}

	return t.Put(rest, val)
}

func (mt *mountBatch) Delete(key ds.Key) error {
	t, rest, err := mt.lookupBatch(key)
	if err != nil {
		return err
	}

	return t.Delete(rest)
}

func (mt *mountBatch) Commit() error {
Łukasz Magiera's avatar
Łukasz Magiera committed
375 376 377
	mt.lk.Lock()
	defer mt.lk.Unlock()

378 379 380 381 382 383 384 385
	for _, t := range mt.mounts {
		err := t.Commit()
		if err != nil {
			return err
		}
	}
	return nil
}
386 387 388 389 390

func (d *Datastore) Check() error {
	for _, m := range d.mounts {
		if c, ok := m.Datastore.(ds.CheckedDatastore); ok {
			if err := c.Check(); err != nil {
391
				return xerrors.Errorf("checking datastore at %s: %w", m.Prefix.String(), err)
392 393 394 395 396 397 398 399 400 401
			}
		}
	}
	return nil
}

func (d *Datastore) Scrub() error {
	for _, m := range d.mounts {
		if c, ok := m.Datastore.(ds.ScrubbedDatastore); ok {
			if err := c.Scrub(); err != nil {
402
				return xerrors.Errorf("scrubbing datastore at %s: %w", m.Prefix.String(), err)
403 404 405 406 407 408 409 410 411 412
			}
		}
	}
	return nil
}

func (d *Datastore) CollectGarbage() error {
	for _, m := range d.mounts {
		if c, ok := m.Datastore.(ds.GCDatastore); ok {
			if err := c.CollectGarbage(); err != nil {
413
				return xerrors.Errorf("gc on datastore at %s: %w", m.Prefix.String(), err)
414 415 416 417 418
			}
		}
	}
	return nil
}