Exemplo n.º 1
0
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
}
Exemplo n.º 2
0
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
}
Exemplo n.º 3
0
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))
	})
}
Exemplo n.º 4
0
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
}
Exemplo n.º 5
0
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))
	})
}
Exemplo n.º 6
0
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
}
Exemplo n.º 7
0
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
}
Exemplo n.º 8
0
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
}