<commit_before>package controller

import (
	"errors"
	"github.com/crawlerclub/x/crawler"
	"github.com/crawlerclub/x/store"
	"github.com/crawlerclub/x/types"
	"github.com/golang/glog"
	"github.com/syndtr/goleveldb/leveldb/util"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"
)

var (
	ErrNotInited      = errors.New("controller/controller.go call Controller.Init first")
	ErrWorkerCount    = errors.New("controller/controller.go worker count must between 0 and 1000")
	ErrNilCrawlerItem = errors.New("controller/controller.go CrawlerItem is nil")
	ErrNamesNotSame   = errors.New("controller/controller.go CrawlerNames are not the same")
	ErrNameTaken      = errors.New("controller/controller.go CrawlerName already taken")
	ErrNoName         = errors.New("controller/controller.go no CrawlerName")
)

var StoreNames = []string{"crawler", "seed", "running", "crontab"}

type Controller struct {
	Crawlers    map[string]crawler.Crawler
	Schduler    CrawlerScheduler
	Stores      map[string]*store.LevelStore
	WorkerCount int

	workDir  string
	isInited bool
}

func timeStr(t int64) string {
	return time.Unix(t, 0).Format("20060102030405")
}

func (self *Controller) Init(dir string, wc int) error {
	glog.Info("call Init")
	if wc <= 0 || wc > 1000 {
		return ErrWorkerCount
	}
	self.workDir = dir
	self.WorkerCount = wc
	_, err := os.Stat(dir)
	if os.IsNotExist(err) {
		err = os.MkdirAll(dir, 0755)
		if err != nil {
			return err
		}
	}
	self.Stores = make(map[string]*store.LevelStore)
	for _, name := range StoreNames {
		self.Stores[name], err = store.NewLevelStore(dir + "/db/" + name)
		if err != nil {
			return err
		}
	}

	err = self.initCrawlersFromDB()
	if err != nil {
		return err
	}
	self.isInited = true
	return nil
}

func (self *Controller) runCrawler(item *types.CrawlerItem) error {
	glog.Info("call runCrawler: ", item.CrawlerName)
	var crawler crawler.Crawler
	crawler.Conf = &item.Conf
	err := crawler.InitEs()
	if err != nil {
		//return err
		glog.Error(err)
	}
	dir := self.workDir + "/queue"
	err = crawler.InitTaskQueue(dir)
	if err != nil {
		return err
	}
	if _, ok := self.Crawlers[item.CrawlerName]; ok {
		self.Schduler.Remove(item.CrawlerName)
	}
	self.Crawlers[item.CrawlerName] = crawler
	var sitem Item
	sitem.CrawlerName = item.CrawlerName
	sitem.Weight = item.Weight
	glog.Info("SortedList add ", sitem)
	err = self.Schduler.Insert(sitem)
	if err != nil {
		return err
	}
	for _, url := range item.Conf.StartUrls {
		task := types.Task{
			CrawlerName: item.Conf.CrawlerName,
			ParserName:  item.Conf.StartParserName,
			IsSeedUrl:   true,
			Url:         url,
		}
		if ret, _ := self.Stores["seed"].Has(task.Id()); !ret {
			// enqueue new start urls
			if _, err = crawler.TaskQueue.EnqueueObject(task); err != nil {
				return err
			}
		}
		value, _ := store.ObjectToBytes(task)
		self.Stores["seed"].Put(task.Id(), value) // update seed
	}
	return nil
}

func (self *Controller) initCrawlersFromDB() error {
	glog.Info("call initCrawlersFromDB")
	self.Crawlers = make(map[string]crawler.Crawler)
	self.Schduler.Init()

	crawlerStore := self.Stores["crawler"]
	count := 0
	err := crawlerStore.ForEach(nil, func(key, value []byte) (bool, error) {
		count += 1
		var item types.CrawlerItem
		e := store.BytesToObject(value, &item)
		if e != nil {
			return false, e
		}
		e = self.runCrawler(&item)
		if e != nil {
			return false, nil
		}
		return true, nil<commit_msg>fix bug, call runCrawler for enabled items only when initCrawlersFromDB<commit_after>package controller

import (
	"errors"
	"github.com/crawlerclub/x/crawler"
	"github.com/crawlerclub/x/store"
	"github.com/crawlerclub/x/types"
	"github.com/golang/glog"
	"github.com/syndtr/goleveldb/leveldb/util"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"
)

var (
	ErrNotInited      = errors.New("controller/controller.go call Controller.Init first")
	ErrWorkerCount    = errors.New("controller/controller.go worker count must between 0 and 1000")
	ErrNilCrawlerItem = errors.New("controller/controller.go CrawlerItem is nil")
	ErrNamesNotSame   = errors.New("controller/controller.go CrawlerNames are not the same")
	ErrNameTaken      = errors.New("controller/controller.go CrawlerName already taken")
	ErrNoName         = errors.New("controller/controller.go no CrawlerName")
)

var StoreNames = []string{"crawler", "seed", "running", "crontab"}

type Controller struct {
	Crawlers    map[string]crawler.Crawler
	Schduler    CrawlerScheduler
	Stores      map[string]*store.LevelStore
	WorkerCount int

	workDir  string
	isInited bool
}

func timeStr(t int64) string {
	return time.Unix(t, 0).Format("20060102030405")
}

func (self *Controller) Init(dir string, wc int) error {
	glog.Info("call Init")
	if wc <= 0 || wc > 1000 {
		return ErrWorkerCount
	}
	self.workDir = dir
	self.WorkerCount = wc
	_, err := os.Stat(dir)
	if os.IsNotExist(err) {
		err = os.MkdirAll(dir, 0755)
		if err != nil {
			return err
		}
	}
	self.Stores = make(map[string]*store.LevelStore)
	for _, name := range StoreNames {
		self.Stores[name], err = store.NewLevelStore(dir + "/db/" + name)
		if err != nil {
			return err
		}
	}

	err = self.initCrawlersFromDB()
	if err != nil {
		return err
	}
	self.isInited = true
	return nil
}

func (self *Controller) runCrawler(item *types.CrawlerItem) error {
	glog.Info("call runCrawler: ", item.CrawlerName)
	var crawler crawler.Crawler
	crawler.Conf = &item.Conf
	err := crawler.InitEs()
	if err != nil {
		//return err
		glog.Error(err)
	}
	dir := self.workDir + "/queue"
	err = crawler.InitTaskQueue(dir)
	if err != nil {
		return err
	}
	if _, ok := self.Crawlers[item.CrawlerName]; ok {
		self.Schduler.Remove(item.CrawlerName)
	}
	self.Crawlers[item.CrawlerName] = crawler
	var sitem Item
	sitem.CrawlerName = item.CrawlerName
	sitem.Weight = item.Weight
	glog.Info("SortedList add ", sitem)
	err = self.Schduler.Insert(sitem)
	if err != nil {
		return err
	}
	for _, url := range item.Conf.StartUrls {
		task := types.Task{
			CrawlerName: item.Conf.CrawlerName,
			ParserName:  item.Conf.StartParserName,
			IsSeedUrl:   true,
			Url:         url,
		}
		if ret, _ := self.Stores["seed"].Has(task.Id()); !ret {
			// enqueue new start urls
			if _, err = crawler.TaskQueue.EnqueueObject(task); err != nil {
				return err
			}
		}
		value, _ := store.ObjectToBytes(task)
		self.Stores["seed"].Put(task.Id(), value) // update seed
	}
	return nil
}

func (self *Controller) initCrawlersFromDB() error {
	glog.Info("call initCrawlersFromDB")
	self.Crawlers = make(map[string]crawler.Crawler)
	self.Schduler.Init()

	crawlerStore := self.Stores["crawler"]
	count := 0
	err := crawlerStore.ForEach(nil, func(key, value []byte) (bool, error) {
		count += 1
		var item types.CrawlerItem
		e := store.BytesToObject(value, &item)
		if e != nil {
			return false, e
		}
		if item.Status == "enabled" {
			e = self.runCrawler(&item)
			if e != nil {
				return false, nil
			}
		}
		return true, nil
{'fraction_non_alphanumeric': 0.13410404624277455, 'fraction_numerical': 0.00953757225433526, 'mean_word_length': 3.7371663244353184, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5182648', 'n_tokens_mistral': 2795, 'n_tokens_neox': 2606, 'n_words': 565}