func New(path string, opt *bolt.Options, e codec.Codec) (bs *BoltStore, err error) {
if e == nil {
e = codec.JSON
}
bs = &BoltStore{
filter: bloom.NewFilter(-1, 0.0001),
codec: e,
}
if bs.DB, err = bolt.Open(path, 0644, opt); err != nil {
return nil, err
}
err = bs.DB.Update(func(tx *bolt.Tx) error {
if _, err = tx.CreateBucketIfNotExists(bkURL); err != nil {
return err
}
b, err := tx.CreateBucketIfNotExists(bkCount)
if err != nil {
return err
}
for _, k := range [][]byte{
keyVisitCount, keyURLCount, keyErrorCount, keyFinishCount,
} {
if _, err := bkPutNX(b, k, util.I64tob(0)); err != nil {
return err
}
}
return nil
})
if err != nil {
bs = nil
}
return
}
func New(path string, o *opt.Options, e codec.Codec) (s *LevelStore, err error) {
db, err := leveldb.OpenFile(path, o)
if err != nil {
return
}
for _, k := range [][]byte{
keyVisitCount, keyURLCount, keyErrorCount, keyFinishCount,
} {
var has bool
if has, err = db.Has(k, nil); err != nil {
return
} else if has {
continue
}
if err = db.Put(k, util.I64tob(0), nil); err != nil {
return
}
}
if e == nil {
e = codec.JSON
}
return &LevelStore{
DB: db,
codec: e,
}, nil
}
func (s *BoltStore) IncErrorCount() (err error) {
return s.DB.Update(func(tx *bolt.Tx) error {
b := tx.Bucket(bkCount)
cnt := util.Btoi64(b.Get(keyErrorCount)) + 1
return b.Put(keyErrorCount, util.I64tob(cnt))
})
}
func (s *LevelStore) incCount(k []byte) (err error) {
tx, err := s.DB.OpenTransaction()
if err != nil {
return
}
commit := false
var v []byte
if v, err = tx.Get(k, nil); err == nil {
cnt := util.Btoi64(v) + 1
if err = tx.Put(k, util.I64tob(cnt), nil); err == nil {
commit = true
err = tx.Commit()
}
}
if !commit && err != nil {
tx.Discard() // TODO: handle error
}
return
}
func (s *BoltStore) Complete(u *url.URL) error {
return s.DB.Update(func(tx *bolt.Tx) error {
b := tx.Bucket(bkURL)
uu, err := s.getFromBucket(b, u)
if err != nil {
return err
}
uu.Done = true
k := []byte(u.String())
w := &wrapper{}
v, err := s.codec.Marshal(w.From(uu))
if err != nil {
return err
}
if err = b.Put(k, v); err != nil {
return err
}
b = tx.Bucket(bkCount)
cnt := util.Btoi64(b.Get(keyFinishCount)) + 1
return b.Put(keyFinishCount, util.I64tob(cnt))
})
}
func (s *LevelStore) PutNX(u *crawler.URL) (ok bool, err error) {
tx, err := s.DB.OpenTransaction()
if err != nil {
return
}
commit := false
defer func() {
if !commit && (err != nil || !ok) {
tx.Discard() // TODO: handle error
}
}()
key := keyURL(&u.URL)
has, err := tx.Has(key, nil)
if err != nil {
return
} else if has {
return false, nil
}
v, err := s.codec.Marshal(u)
if err != nil {
return
}
if err = tx.Put(key, v, nil); err != nil {
return
}
if v, err = tx.Get(keyURLCount, nil); err == nil {
cnt := util.Btoi64(v) + 1
if err = tx.Put(keyURLCount, util.I64tob(cnt), nil); err == nil {
commit = true
if err = tx.Commit(); err == nil {
ok = true
}
}
}
return
}
func (s *LevelStore) Complete(u *url.URL) (err error) {
tx, err := s.DB.OpenTransaction()
if err != nil {
return
}
commit := false
defer func() {
if !commit && err != nil {
tx.Discard() // TODO: handle error
}
}()
key := keyURL(u)
v, err := tx.Get(key, nil)
if err != nil {
return
}
uu := crawler.URL{}
if err = s.codec.Unmarshal(v, &uu); err != nil {
return
}
uu.Done = true
if v, err = s.codec.Marshal(&uu); err != nil {
return
}
if err = tx.Put(key, v, nil); err != nil {
return
}
if v, err = tx.Get(keyFinishCount, nil); err != nil {
return
}
cnt := util.Btoi64(v) + 1
if err = tx.Put(keyFinishCount, util.I64tob(cnt), nil); err == nil {
commit = true
err = tx.Commit()
}
return
}
func (s *BoltStore) PutNX(u *crawler.URL) (ok bool, err error) {
err = s.DB.Update(func(tx *bolt.Tx) error {
b := tx.Bucket(bkURL)
k := []byte(u.URL.String())
w := &wrapper{}
v, err := s.codec.Marshal(w.From(u))
if err != nil {
return err
}
if ok, err = bkPutNX(b, k, v); err != nil {
return err
} else if !ok {
return nil
}
b = tx.Bucket(bkCount)
cnt := util.Btoi64(b.Get(keyURLCount)) + 1
return b.Put(keyURLCount, util.I64tob(cnt))
})
if err == nil && ok {
s.filter.Add(&u.URL)
}
return
}
