openbilibili/vendor/github.com/go-ego/riot/riot.go

// Copyright 2017 ego authors
//
// Licensed under the Apache License, Version 2.0 (the "License"): you may
// not use this file except in compliance with the License. You may obtain
// a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
// License for the specific language governing permissions and limitations
// under the License.

package riot

import (
	"bytes"
	"fmt"
	"log"
	"os"
	"strings"

	"encoding/binary"
	"encoding/gob"

	"github.com/go-ego/murmur"
	"github.com/go-ego/riot/core"
	"github.com/go-ego/riot/types"
	toml "github.com/go-vgo/gt/conf"
)

// New create a new engine with mode
func New(conf ...interface{}) *Engine {
	// func (engine *Engine) New(conf com.Config) *Engine{
	if len(conf) > 0 && strings.HasSuffix(conf[0].(string), ".toml") {
		var (
			config   types.EngineOpts
			searcher = &Engine{}
		)

		fs := conf[0].(string)
		log.Println("conf path is: ", fs)
		toml.Init(fs, &config)
		go toml.Watch(fs, &config)

		searcher.Init(config)
		return searcher
	}

	return NewEngine(conf...)
}

// NewEngine create a new engine
func NewEngine(conf ...interface{}) *Engine {
	var (
		searcher = &Engine{}

		path          = DefaultPath
		storageShards = 10
		numShards     = 10

		segmentDict string
	)

	if len(conf) > 0 {
		segmentDict = conf[0].(string)
	}

	if len(conf) > 1 {
		path = conf[1].(string)
	}

	if len(conf) > 2 {
		numShards = conf[2].(int)
		storageShards = conf[2].(int)
	}

	searcher.Init(types.EngineOpts{
		// Using:         using,
		StoreShards: storageShards,
		NumShards:   numShards,
		IndexerOpts: &types.IndexerOpts{
			IndexType: types.DocIdsIndex,
		},
		UseStore:    true,
		StoreFolder: path,
		// StoreEngine: storageEngine,
		GseDict: segmentDict,
		// StopTokenFile: stopTokenFile,
	})

	// defer searcher.Close()
	os.MkdirAll(path, 0777)

	// 等待索引刷新完毕
	// searcher.Flush()
	// log.Println("recover index number: ", searcher.NumDocsIndexed())

	return searcher
}

// func (engine *Engine) IsDocExist(docId uint64) bool {
// 	return core.IsDocExist(docId)
// }

// HasDoc if the document is exist return true
func (engine *Engine) HasDoc(docId uint64) bool {
	for shard := 0; shard < engine.initOptions.NumShards; shard++ {
		engine.indexers = append(engine.indexers, core.Indexer{})

		has := engine.indexers[shard].HasDoc(docId)

		if has {
			return true
		}
	}

	return false
}

// HasDocDB if the document is exist in the database
// return true
func (engine *Engine) HasDocDB(docId uint64) bool {
	b := make([]byte, 10)
	length := binary.PutUvarint(b, docId)

	shard := murmur.Sum32(fmt.Sprintf("%d", docId)) %
		uint32(engine.initOptions.StoreShards)

	has, err := engine.dbs[shard].Has(b[0:length])
	if err != nil {
		log.Println("engine.dbs[shard].Has(b[0:length]): ", err)
	}

	return has
}

// GetDBAllIds get all the DocId from the storage database
// and return
// 从数据库遍历所有的 DocId, 并返回
func (engine *Engine) GetDBAllIds() []uint64 {
	docsId := make([]uint64, 0)
	for i := range engine.dbs {
		engine.dbs[i].ForEach(func(k, v []byte) error {
			// fmt.Println(k, v)
			docId, _ := binary.Uvarint(k)
			docsId = append(docsId, docId)
			return nil
		})
	}

	return docsId
}

// GetDBAllDocs get the db all docs
func (engine *Engine) GetDBAllDocs() (
	docsId []uint64, docsData []types.DocData) {
	for i := range engine.dbs {
		engine.dbs[i].ForEach(func(key, val []byte) error {
			// fmt.Println(k, v)
			docId, _ := binary.Uvarint(key)
			docsId = append(docsId, docId)

			buf := bytes.NewReader(val)
			dec := gob.NewDecoder(buf)

			var data types.DocData
			err := dec.Decode(&data)
			if err != nil {
				log.Println("dec.decode: ", err)
			}

			docsData = append(docsData, data)

			return nil
		})
	}

	return docsId, docsData
}

// GetAllDocIds get all the DocId from the storage database
// and return
// 从数据库遍历所有的 DocId, 并返回
func (engine *Engine) GetAllDocIds() []uint64 {
	return engine.GetDBAllIds()
}

// Try handler(err)
func Try(fun func(), handler func(interface{})) {
	defer func() {
		if err := recover(); err != nil {
			handler(err)
		}
	}()
	fun()
}