func (s *BoltStore) IsFinished() (is bool, err error) {
err = s.DB.View(func(tx *bolt.Tx) error {
b := tx.Bucket(bkCount)
finish := util.Btoi64(b.Get(keyFinishCount))
urlcnt := util.Btoi64(b.Get(keyURLCount))
is = finish >= urlcnt
return nil
})
return
}
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) IsFinished() (is bool, err error) {
snap, err := s.DB.GetSnapshot()
if err != nil {
return
}
defer snap.Release()
v, err := snap.Get(keyURLCount, nil)
if err != nil {
return
}
urlcnt := util.Btoi64(v)
if v, err = snap.Get(keyFinishCount, nil); err == nil {
if finish := util.Btoi64(v); finish >= urlcnt {
is = true
}
}
return
}
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
}