1	/*
     2	Copyright 2011 The Perkeep Authors
     3	
     4	Licensed under the Apache License, Version 2.0 (the "License");
     5	you may not use this file except in compliance with the License.
     6	You may obtain a copy of the License at
     7	
     8	     http://www.apache.org/licenses/LICENSE-2.0
     9	
    10	Unless required by applicable law or agreed to in writing, software
    11	distributed under the License is distributed on an "AS IS" BASIS,
    12	WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13	See the License for the specific language governing permissions and
    14	limitations under the License.
    15	*/
    16	
    17	package sorted
    18	
    19	import (
    20		"errors"
    21		"log"
    22		"sync"
    23	
    24		"github.com/syndtr/goleveldb/leveldb/comparer"
    25		"github.com/syndtr/goleveldb/leveldb/iterator"
    26		"github.com/syndtr/goleveldb/leveldb/memdb"
    27		"github.com/syndtr/goleveldb/leveldb/util"
    28		"go4.org/jsonconfig"
    29	)
    30	
    31	// NewMemoryKeyValue returns a KeyValue implementation that's backed only
    32	// by memory. It's mostly useful for tests and development.
    33	func NewMemoryKeyValue() KeyValue {
    34		kv := new(memKeys)
    35		kv.Wipe()
    36		return kv
    37	}
    38	
    39	var _ Wiper = (*memKeys)(nil)
    40	
    41	// memKeys is a naive in-memory implementation of KeyValue for test & development
    42	// purposes only.
    43	type memKeys struct {
    44		mu sync.Mutex // guards db
    45		db *memdb.DB
    46	}
    47	
    48	func (s *memKeys) Wipe() error {
    49		s.mu.Lock()
    50		defer s.mu.Unlock()
    51		s.db = memdb.New(comparer.DefaultComparer, 128)
    52		return nil
    53	}
    54	
    55	// memIter converts from leveldb's iterator.Iterator interface, which
    56	// operates on []byte, to Perkeep's index.Iterator, which operates
    57	// on string.
    58	type memIter struct {
    59		lit  iterator.Iterator // underlying leveldb iterator
    60		k, v *string           // if nil, not stringified yet
    61	}
    62	
    63	func (s *memIter) Next() bool {
    64		s.k, s.v = nil, nil
    65		return s.lit.Next()
    66	}
    67	
    68	func (s *memIter) Close() error {
    69		s.lit.Release()
    70		return s.lit.Error()
    71	}
    72	
    73	func (s *memIter) KeyBytes() []byte {
    74		return s.lit.Key()
    75	}
    76	
    77	func (s *memIter) ValueBytes() []byte {
    78		return s.lit.Value()
    79	}
    80	
    81	func (s *memIter) Key() string {
    82		if s.k != nil {
    83			return *s.k
    84		}
    85		str := string(s.KeyBytes())
    86		s.k = &str
    87		return str
    88	}
    89	
    90	func (s *memIter) Value() string {
    91		if s.v != nil {
    92			return *s.v
    93		}
    94		str := string(s.ValueBytes())
    95		s.v = &str
    96		return str
    97	}
    98	
    99	func (mk *memKeys) Get(key string) (string, error) {
   100		mk.mu.Lock()
   101		defer mk.mu.Unlock()
   102		k, err := mk.db.Get([]byte(key))
   103		if err == memdb.ErrNotFound {
   104			return "", ErrNotFound
   105		}
   106		return string(k), err
   107	}
   108	
   109	func (mk *memKeys) Find(start, end string) Iterator {
   110		mk.mu.Lock()
   111		defer mk.mu.Unlock()
   112		var startB, endB []byte
   113		if start != "" {
   114			startB = []byte(start)
   115		}
   116		if end != "" {
   117			endB = []byte(end)
   118		}
   119		lit := mk.db.NewIterator(&util.Range{Start: startB, Limit: endB})
   120		return &memIter{lit: lit}
   121	}
   122	
   123	func (mk *memKeys) Set(key, value string) error {
   124		if err := CheckSizes(key, value); err != nil {
   125			log.Printf("Skipping storing (%q:%q): %v", key, value, err)
   126			return nil
   127		}
   128		mk.mu.Lock()
   129		defer mk.mu.Unlock()
   130		return mk.db.Put([]byte(key), []byte(value))
   131	}
   132	
   133	func (mk *memKeys) Delete(key string) error {
   134		mk.mu.Lock()
   135		defer mk.mu.Unlock()
   136		err := mk.db.Delete([]byte(key))
   137		if err == memdb.ErrNotFound {
   138			return nil
   139		}
   140		return err
   141	}
   142	
   143	func (mk *memKeys) BeginBatch() BatchMutation {
   144		return &batch{}
   145	}
   146	
   147	func (mk *memKeys) CommitBatch(bm BatchMutation) error {
   148		b, ok := bm.(*batch)
   149		if !ok {
   150			return errors.New("invalid batch type; not an instance returned by BeginBatch")
   151		}
   152		mk.mu.Lock()
   153		defer mk.mu.Unlock()
   154		for _, m := range b.Mutations() {
   155			if m.IsDelete() {
   156				if err := mk.db.Delete([]byte(m.Key())); err != nil && err != memdb.ErrNotFound {
   157					return err
   158				}
   159			} else {
   160				// TODO(mpl): we need to force a reindex when we have a proper solution
   161				// for storing these too large attributes, if ever.
   162				if err := CheckSizes(m.Key(), m.Value()); err != nil {
   163					log.Printf("Skipping storing (%q:%q): %v", m.Key(), m.Value(), err)
   164					continue
   165				}
   166				if err := mk.db.Put([]byte(m.Key()), []byte(m.Value())); err != nil {
   167					return err
   168				}
   169			}
   170		}
   171		return nil
   172	}
   173	
   174	func (mk *memKeys) Close() error { return nil }
   175	
   176	func init() {
   177		RegisterKeyValue("memory", func(cfg jsonconfig.Obj) (KeyValue, error) {
   178			if err := cfg.Validate(); err != nil {
   179				return nil, err
   180			}
   181			return NewMemoryKeyValue(), nil
   182		})
   183	}