github.com/amazechain/amc@v0.1.3/internal/amcdb/lmdb/db_rw.go (about)

     1  // Copyright 2022 The AmazeChain Authors
     2  // This file is part of the AmazeChain library.
     3  //
     4  // The AmazeChain library is free software: you can redistribute it and/or modify
     5  // it under the terms of the GNU Lesser General Public License as published by
     6  // the Free Software Foundation, either version 3 of the License, or
     7  // (at your option) any later version.
     8  //
     9  // The AmazeChain library is distributed in the hope that it will be useful,
    10  // but WITHOUT ANY WARRANTY; without even the implied warranty of
    11  // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
    12  // GNU Lesser General Public License for more details.
    13  //
    14  // You should have received a copy of the GNU Lesser General Public License
    15  // along with the AmazeChain library. If not, see <http://www.gnu.org/licenses/>.
    16  
    17  package lmdb
    18  
    19  import (
    20  	"context"
    21  	"fmt"
    22  	"github.com/amazechain/amc/common/db"
    23  	"github.com/erigontech/mdbx-go/mdbx"
    24  	"sync"
    25  )
    26  
    27  type DBI struct {
    28  	*mdbx.DBI
    29  	env    *mdbx.Env
    30  	ctx    context.Context
    31  	cancel context.CancelFunc
    32  
    33  	lock sync.RWMutex
    34  	name string
    35  
    36  	//log *zap.SugaredLogger
    37  }
    38  
    39  func newDBI(ctx context.Context, env *mdbx.Env, name string) (*DBI, error) {
    40  	var (
    41  		mdb mdbx.DBI
    42  		err error
    43  	)
    44  	err = env.Update(func(txn *mdbx.Txn) error {
    45  		mdb, err = txn.OpenDBI(name, mdbx.Create, nil, nil)
    46  		return err
    47  	})
    48  
    49  	if err != nil {
    50  		return nil, err
    51  	}
    52  	c, cancel := context.WithCancel(ctx)
    53  	dbi := DBI{
    54  		DBI:    &mdb,
    55  		env:    env,
    56  		ctx:    c,
    57  		cancel: cancel,
    58  		name:   name,
    59  	}
    60  
    61  	return &dbi, nil
    62  }
    63  
    64  func (db *DBI) Get(key []byte) (value []byte, err error) {
    65  	db.lock.RLock()
    66  	defer db.lock.RUnlock()
    67  	err = db.env.View(func(txn *mdbx.Txn) error {
    68  		value, err = txn.Get(*db.DBI, key)
    69  		return err
    70  	})
    71  
    72  	return value, err
    73  }
    74  
    75  func (db *DBI) Gets(key []byte, count uint) (keys [][]byte, values [][]byte, err error) {
    76  	db.lock.RLock()
    77  	defer db.lock.RUnlock()
    78  	err = db.env.View(func(txn *mdbx.Txn) error {
    79  		cur, err := txn.OpenCursor(*db.DBI)
    80  		if err != nil {
    81  			return err
    82  		}
    83  		defer cur.Close()
    84  
    85  		k, v, err := cur.Get(key, nil, mdbx.Set) //SetRange
    86  		if err != nil {
    87  			return err
    88  		}
    89  
    90  		keys = append(keys, k)
    91  		values = append(values, v)
    92  		for {
    93  			k, v, err = cur.Get(nil, nil, mdbx.Next)
    94  			if mdbx.IsNotFound(err) {
    95  				break
    96  			}
    97  			if err != nil {
    98  				return err
    99  			}
   100  
   101  			keys = append(keys, k)
   102  			values = append(values, v)
   103  			if len(keys) >= int(count) {
   104  				break
   105  			}
   106  		}
   107  
   108  		return nil
   109  	})
   110  
   111  	return keys, values, err
   112  }
   113  
   114  func (db *DBI) GetIterator(key []byte) (iterator db.IIterator, err error) {
   115  	//panic(fmt.Errorf("don`t support Iterator current"))
   116  	return newIterator(db, key)
   117  }
   118  
   119  func (db *DBI) Put(key []byte, value []byte) (err error) {
   120  	err = db.env.Update(func(txn *mdbx.Txn) error {
   121  		return txn.Put(*db.DBI, key, value, 0) //mdbx.NoOverwrite
   122  	})
   123  
   124  	return err
   125  }
   126  
   127  func (db *DBI) Puts(keys [][]byte, values [][]byte) (err error) {
   128  	db.lock.Lock()
   129  	defer db.lock.Unlock()
   130  	if len(keys) != len(values) {
   131  		return fmt.Errorf("failed the number of keys and values is inconsistent")
   132  	}
   133  	err = db.env.Update(func(txn *mdbx.Txn) error {
   134  		for i := 0; i < len(keys); i++ {
   135  			err = txn.Put(*db.DBI, keys[i], values[i], 0)
   136  			if err != nil {
   137  				return err
   138  			}
   139  		}
   140  		return nil
   141  	})
   142  
   143  	return err
   144  }
   145  
   146  func (db *DBI) Drop() (err error) {
   147  	db.lock.Lock()
   148  	defer db.lock.Unlock()
   149  	return db.env.Update(func(txn *mdbx.Txn) error {
   150  		return txn.Drop(*db.DBI, true)
   151  	})
   152  }
   153  
   154  func (db *DBI) Delete(key []byte) (err error) {
   155  	db.lock.Lock()
   156  	defer db.lock.Unlock()
   157  	return db.env.Update(func(txn *mdbx.Txn) error {
   158  		return txn.Del(*db.DBI, key, nil)
   159  	})
   160  }