itdoxy-lab/vendor/gitea.com/lunny/levelqueue/uniquequeue.go

// Copyright 2020 Andrew Thornton. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package levelqueue

import (
	"fmt"

	"github.com/syndtr/goleveldb/leveldb"
	"github.com/syndtr/goleveldb/leveldb/errors"
)

const (
	uniqueQueuePrefixStr = "unique"
)

// UniqueQueue defines an unique queue struct
type UniqueQueue struct {
	q                 *Queue
	set               *Set
	db                *leveldb.DB
	closeUnderlyingDB bool
}

// OpenUnique opens an unique queue from the db path or creates a set if it doesn't exist.
// The keys in the queue portion will not be prefixed, and the set keys will be prefixed with "set-"
func OpenUnique(dataDir string) (*UniqueQueue, error) {
	db, err := leveldb.OpenFile(dataDir, nil)
	if err != nil {
		if !errors.IsCorrupted(err) {
			return nil, err
		}
		db, err = leveldb.RecoverFile(dataDir, nil)
		if err != nil {
			return nil, err
		}
	}
	return NewUniqueQueue(db, []byte{}, []byte(uniqueQueuePrefixStr), true)
}

// NewUniqueQueue creates a new unique queue from a db.
// The queue keys will be prefixed with queuePrefix and the set keys with setPrefix
// and at close the db will be closed as per closeUnderlyingDB
func NewUniqueQueue(db *leveldb.DB, queuePrefix []byte, setPrefix []byte, closeUnderlyingDB bool) (*UniqueQueue, error) {
	internal, err := NewQueue(db, queuePrefix, false)
	if err != nil {
		return nil, err
	}
	set, err := NewSet(db, setPrefix, false)
	if err != nil {
		return nil, err
	}
	queue := &UniqueQueue{
		q:                 internal,
		set:               set,
		db:                db,
		closeUnderlyingDB: closeUnderlyingDB,
	}

	return queue, err
}

// LPush pushes data to the left of the queue
func (queue *UniqueQueue) LPush(data []byte) error {
	return queue.LPushFunc(data, nil)
}

// LPushFunc pushes data to the left of the queue and calls the callback if it is added
func (queue *UniqueQueue) LPushFunc(data []byte, fn func() error) error {
	added, err := queue.set.Add(data)
	if err != nil {
		return err
	}
	if !added {
		return ErrAlreadyInQueue
	}

	if fn != nil {
		err = fn()
		if err != nil {
			_, remErr := queue.set.Remove(data)
			if remErr != nil {
				return fmt.Errorf("%v & %v", err, remErr)
			}
			return err
		}
	}

	return queue.q.LPush(data)
}

// RPush pushes data to the right of the queue
func (queue *UniqueQueue) RPush(data []byte) error {
	return queue.RPushFunc(data, nil)
}

// RPushFunc pushes data to the right of the queue and calls the callback if is added
func (queue *UniqueQueue) RPushFunc(data []byte, fn func() error) error {
	added, err := queue.set.Add(data)
	if err != nil {
		return err
	}
	if !added {
		return ErrAlreadyInQueue
	}

	if fn != nil {
		err = fn()
		if err != nil {
			_, remErr := queue.set.Remove(data)
			if remErr != nil {
				return fmt.Errorf("%v & %v", err, remErr)
			}
			return err
		}
	}

	return queue.q.RPush(data)
}

// RPop pop data from the right of the queue
func (queue *UniqueQueue) RPop() ([]byte, error) {
	popped, err := queue.q.RPop()
	if err != nil {
		return popped, err
	}
	_, err = queue.set.Remove(popped)

	return popped, err
}

// RHandle receives a user callback function to handle the right element of the queue, if the function returns nil, then delete the element, otherwise keep the element.
func (queue *UniqueQueue) RHandle(h func([]byte) error) error {
	return queue.q.RHandle(func(data []byte) error {
		err := h(data)
		if err != nil {
			return err
		}
		_, err = queue.set.Remove(data)
		return err
	})
}

// LPop pops data from left of the queue
func (queue *UniqueQueue) LPop() ([]byte, error) {
	popped, err := queue.q.LPop()
	if err != nil {
		return popped, err
	}
	_, err = queue.set.Remove(popped)

	return popped, err
}

// LHandle receives a user callback function to handle the left element of the queue, if the function returns nil, then delete the element, otherwise keep the element.
func (queue *UniqueQueue) LHandle(h func([]byte) error) error {
	return queue.q.LHandle(func(data []byte) error {
		err := h(data)
		if err != nil {
			return err
		}
		_, err = queue.set.Remove(data)
		return err
	})
}

// Has checks whether the data is already in the queue
func (queue *UniqueQueue) Has(data []byte) (bool, error) {
	return queue.set.Has(data)
}

// Len returns the length of the queue
func (queue *UniqueQueue) Len() int64 {
	queue.set.lock.Lock()
	defer queue.set.lock.Unlock()
	return queue.q.Len()
}

// Close closes the queue (and the underlying DB if set to closeUnderlyingDB)
func (queue *UniqueQueue) Close() error {
	_ = queue.q.Close()
	_ = queue.set.Close()
	if !queue.closeUnderlyingDB {
		queue.db = nil
		return nil
	}
	err := queue.db.Close()
	queue.db = nil
	return err
}