/store.go

  1package diskv
  2
  3import (
  4	"fmt"
  5	"strings"
  6	"sync"
  7	"os"
  8	"time"
  9	"log"
 10	"io/ioutil"
 11	"path/filepath"
 12	"github.com/petar/GoLLRB/llrb"
 13)
 14
 15type DStore struct {
 16	basedir    string
 17	xf         TransformFunc
 18	fileperm   uint32
 19	dirperm    uint32
 20	cache      map[Key]Value
 21	mutex      *sync.RWMutex
 22	cachesz    uint
 23	cachemaxsz uint
 24	index      *llrb.Tree
 25	lf         llrb.LessFunc
 26}
 27
 28// Returns a new, unordered Store that represents the data in the given
 29// basedir directory. The cache will grow to a maximum maxsz bytes.
 30func NewDStore(basedir string, xf TransformFunc, maxsz uint) *DStore {
 31	s := &DStore{}
 32	s.basedir = fmt.Sprintf("%s/%s", basedir, "diskv")
 33	s.xf = xf
 34	s.fileperm, s.dirperm = 0777, 0777
 35	s.cache = make(map[Key]Value)
 36	s.mutex = &sync.RWMutex{}
 37	s.cachesz, s.cachemaxsz = 0, maxsz
 38	s.index, s.lf = nil, nil
 39	return s
 40}
 41
 42// Returns a new, ordered Store that represents the data in the given
 43// basedir directory. The directory tree at basedir will be scanned at
 44// initialization time, to populate the index with existing Keys.
 45func NewODStore(basedir string, xf TransformFunc, maxsz uint, lf llrb.LessFunc) *DStore {
 46	s := NewDStore(basedir, xf, maxsz)
 47	s.index = llrb.New(lf)
 48	s.lf = lf
 49	s.populate_index()
 50	return s
 51}
 52
 53func (s *DStore) populate_index() {
 54	t1 := time.Nanoseconds() / 1e6
 55	c := s.Keys()
 56	for k := <-c; len(k) > 0; k = <-c {
 57		s.index.ReplaceOrInsert(k)
 58	}
 59	td := (time.Nanoseconds() / 1e6) - t1
 60	log.Printf("index populated with %d elements in %d ms\n", s.index.Len(), td)
 61}
 62
 63// Returns the directory that will hold the given Key.
 64func (s *DStore) dir(k Key) string {
 65	pathlist := s.xf(k)
 66	return fmt.Sprintf("%s/%s", s.basedir, strings.Join(pathlist, "/"))
 67}
 68
 69// Returns the full path to the file holding data for the given Key.
 70func (s *DStore) filename(k Key) string {
 71	return fmt.Sprintf("%s/%s", s.dir(k), k)
 72}
 73
 74// Creates all necessary directories on disk to hold the given Key.
 75func (s *DStore) ensure_dir(k Key) os.Error {
 76	return os.MkdirAll(s.dir(k), s.dirperm)
 77}
 78
 79// Deletes empty directories in the path walk leading to the Key k.
 80// Typically this function is called after an Erase() is made.
 81func (s *DStore) prune_dirs(k Key) os.Error {
 82	pathlist := s.xf(k)
 83	for i := range pathlist {
 84		pslice := pathlist[:len(pathlist)-i]
 85		dir := fmt.Sprintf("%s/%s", s.basedir, strings.Join(pslice, "/"))
 86		// thanks to Steven Blenkinsop for this snippet
 87		switch fi, err := os.Stat(dir); true {
 88		case err != nil:
 89			return err
 90		case !fi.IsDirectory():
 91			panic(fmt.Sprintf("corrupt dirstate at %s", dir))
 92		case fi.Nlink != 2:
 93			return nil
 94		default:
 95			if err = os.Remove(dir); err != nil {
 96				return err
 97			}
 98		}
 99	}
100	return nil
101}
102
103// Deletes entries from the cache until it has at least sz bytes available.
104func (s *DStore) ensure_space(sz uint) {
105	for k, v := range s.cache {
106		if (s.cachesz + sz) <= s.cachemaxsz {
107			break
108		}
109		s.cachesz -= uint(len(v))   // len should return uint :|
110		s.cache[k] = Value{}, false // delete is safe, per spec
111	}
112}
113
114func (s *DStore) cache_with_lock(k Key, v Value) {
115	vsz := uint(len(v))
116	s.ensure_space(vsz)
117	if (s.cachesz + vsz) <= s.cachemaxsz {
118		s.cache[k] = v
119		s.cachesz += vsz
120	}
121	if s.index != nil {
122		s.index.ReplaceOrInsert(k)
123	}
124}
125
126func (s *DStore) cache_without_lock(k Key, v Value) {
127	s.mutex.Lock()
128	defer s.mutex.Unlock()
129	s.cache_with_lock(k, v)
130}
131
132// Writes the Value to the store under the given Key.
133// Overwrites anything already existing under the given Key.
134func (s *DStore) Write(k Key, v Value) os.Error {
135	s.mutex.Lock()
136	defer s.mutex.Unlock()
137	if len(k) <= 0 {
138		return os.NewError("empty key")
139	}
140	// write the data to disk (overwrite if exists)
141	if err := s.ensure_dir(k); err != nil {
142		return err
143	}
144	mode := os.O_WRONLY | os.O_CREATE
145	if f, err := os.OpenFile(s.filename(k), mode, s.fileperm); err == nil {
146		defer f.Close()
147		if _, err = f.Write(v); err != nil {
148			return err
149		}
150	} else {
151		return err
152	}
153	// index immediately
154	if s.index != nil {
155		s.index.ReplaceOrInsert(k)
156	}
157	// cache only on read
158	return nil
159}
160
161// Returns the Value stored under the given Key, if it exists.
162// Returns error if the Key doesn't exist.
163func (s *DStore) Read(k Key) (Value, os.Error) {
164	s.mutex.RLock()
165	defer s.mutex.RUnlock()
166	// check cache first
167	if v, ok := s.cache[k]; ok {
168		return v, nil
169	}
170	// read from disk
171	v, err := ioutil.ReadFile(s.filename(k))
172	if err != nil {
173		return Value{}, err
174	}
175	// cache lazily
176	go s.cache_without_lock(k, v)
177	return v, nil
178}
179
180// Erases the Value stored under the given Key, if it exists.
181// Returns error if the Key doesn't exist.
182func (s *DStore) Erase(k Key) os.Error {
183	s.mutex.Lock()
184	defer s.mutex.Unlock()
185	// erase from cache
186	if v, ok := s.cache[k]; ok {
187		s.cachesz -= uint(len(v))
188		s.cache[k] = Value{}, false
189	}
190	// erase from index
191	if s.index != nil {
192		s.index.Delete(k)
193	}
194	// erase from disk
195	filename := s.filename(k)
196	if s, err := os.Stat(filename); err == nil {
197		if !s.IsRegular() {
198			return os.NewError("bad key")
199		}
200		if err = os.Remove(filename); err != nil {
201			return err
202		}
203	} else {
204		return err
205	}
206	if s.index != nil {
207		s.index.Delete(k)
208	}
209	// clean up
210	s.prune_dirs(k)
211	return nil
212}
213
214// Erase all Keys and Values from the DStore.
215func (s *DStore) Flush() os.Error {
216	s.mutex.Lock()
217	defer s.mutex.Unlock()
218	// flush cache
219	s.cache = make(map[Key]Value)
220	s.cachesz = 0
221	// flush index
222	if s.index != nil {
223		s.index.Init(s.lf)
224	}
225	// flush disk
226	return os.RemoveAll(s.basedir)
227}
228
229type DiskvVisitor struct {
230	c chan<- Key
231}
232
233func (v DiskvVisitor) VisitDir(path string, f *os.FileInfo) bool {
234	return true // descend
235}
236
237func (v DiskvVisitor) VisitFile(path string, f *os.FileInfo) {
238	v.c <- Key(f.Name)
239}
240
241// Returns channel generating list of all keys in the store,
242// based on the state of the disk. No order guarantee.
243func (s *DStore) Keys() <-chan Key {
244	c := make(chan Key)
245	go func() {
246		filepath.Walk(s.basedir, DiskvVisitor{c}, nil)
247		close(c)
248	}()
249	return c
250}
251
252func (s *DStore) IsCached(k Key) bool {
253	s.mutex.Lock()
254	defer s.mutex.Unlock()
255	_, present := s.cache[k]
256	return present
257}
258
259func (s *DStore) IsIndexed(k Key) bool {
260	if s.index != nil {
261		s.mutex.Lock()
262		defer s.mutex.Unlock()
263		return s.index.Has(k)
264	}
265	return false
266}
267
268// Return a maximum of count keys from the index, starting at 
269// one beyond the given key and iterating forward
270func (s *DStore) KeysFrom(k Key, count int) ([]Key, os.Error) {
271	if s.index == nil {
272		panic("KeysFrom cannot be called on non-ordered store")
273	}
274	if s.index.Len() <= 0 {
275		return []Key{}, nil
276	}
277	skip_first := true
278	if len(k) <= 0 || !s.index.Has(k) {
279		k = s.index.Min().(Key) // no such key, so start at the top
280		skip_first = false
281	}
282	keys := make([]Key, count)
283	c := s.index.IterRange(k, s.index.Max())
284	total := 0
285	if skip_first {
286		<-c
287	}
288	for i, k := 0, <-c; i < count && k != nil; i, k = i+1, <-c {
289		keys[i] = k.(Key)
290		total++
291	}
292	if total < count { // hack to get around IterRange returning only E < @upper
293		keys[total] = s.index.Max().(Key)
294		total++
295	}
296	keys = keys[:total]
297	return keys, nil
298}