feat: added modules

core/task/stats/service_v2.go

@@ -0,0 +1,143 @@
+package stats
+
+import (
+	"github.com/crawlab-team/crawlab/core/interfaces"
+	"github.com/crawlab-team/crawlab/core/models/models"
+	"github.com/crawlab-team/crawlab/core/models/service"
+	nodeconfig "github.com/crawlab-team/crawlab/core/node/config"
+	"github.com/crawlab-team/crawlab/core/result"
+	"github.com/crawlab-team/crawlab/core/task/log"
+	"github.com/crawlab-team/go-trace"
+	"go.mongodb.org/mongo-driver/bson"
+	"go.mongodb.org/mongo-driver/bson/primitive"
+	"sync"
+	"time"
+)
+
+type ServiceV2 struct {
+	// dependencies
+	nodeCfgSvc interfaces.NodeConfigService
+	modelSvc   service.ModelService
+
+	// internals
+	mu             sync.Mutex
+	resultServices sync.Map
+	rsTtl          time.Duration
+	logDriver      log.Driver
+}
+
+func (svc *ServiceV2) Init() (err error) {
+	go svc.cleanup()
+	return nil
+}
+
+func (svc *ServiceV2) InsertData(id primitive.ObjectID, records ...interface{}) (err error) {
+	resultSvc, err := svc.getResultService(id)
+	if err != nil {
+		return err
+	}
+	if err := resultSvc.Insert(records...); err != nil {
+		return err
+	}
+	go svc.updateTaskStats(id, len(records))
+	return nil
+}
+
+func (svc *ServiceV2) InsertLogs(id primitive.ObjectID, logs ...string) (err error) {
+	return svc.logDriver.WriteLines(id.Hex(), logs)
+}
+
+func (svc *ServiceV2) getResultService(id primitive.ObjectID) (resultSvc interfaces.ResultService, err error) {
+	// atomic operation
+	svc.mu.Lock()
+	defer svc.mu.Unlock()
+
+	// attempt to get from cache
+	res, _ := svc.resultServices.Load(id.Hex())
+	if res != nil {
+		// hit in cache
+		resultSvc, ok := res.(interfaces.ResultService)
+		resultSvc.SetTime(time.Now())
+		if ok {
+			return resultSvc, nil
+		}
+	}
+
+	// task
+	t, err := svc.modelSvc.GetTaskById(id)
+	if err != nil {
+		return nil, err
+	}
+
+	// result service
+	resultSvc, err = result.GetResultService(t.SpiderId)
+	if err != nil {
+		return nil, err
+	}
+
+	// store in cache
+	svc.resultServices.Store(id.Hex(), resultSvc)
+
+	return resultSvc, nil
+}
+
+func (svc *ServiceV2) updateTaskStats(id primitive.ObjectID, resultCount int) {
+	err := service.NewModelServiceV2[models.TaskStatV2]().UpdateById(id, bson.M{
+		"$inc": bson.M{
+			"result_count": resultCount,
+		},
+	})
+	if err != nil {
+		trace.PrintError(err)
+	}
+}
+
+func (svc *ServiceV2) cleanup() {
+	for {
+		// atomic operation
+		svc.mu.Lock()
+
+		svc.resultServices.Range(func(key, value interface{}) bool {
+			rs := value.(interfaces.ResultService)
+			if time.Now().After(rs.GetTime().Add(svc.rsTtl)) {
+				svc.resultServices.Delete(key)
+			}
+			return true
+		})
+
+		svc.mu.Unlock()
+
+		time.Sleep(10 * time.Minute)
+	}
+}
+
+func NewTaskStatsServiceV2() (svc2 *ServiceV2, err error) {
+	// service
+	svc := &ServiceV2{
+		mu:             sync.Mutex{},
+		resultServices: sync.Map{},
+	}
+
+	svc.nodeCfgSvc = nodeconfig.GetNodeConfigService()
+
+	// log driver
+	svc.logDriver, err = log.GetLogDriver(log.DriverTypeFile)
+	if err != nil {
+		return nil, err
+	}
+
+	return svc, nil
+}
+
+var _serviceV2 *ServiceV2
+
+func GetTaskStatsServiceV2() (svr *ServiceV2, err error) {
+	if _serviceV2 != nil {
+		return _serviceV2, nil
+	}
+	_serviceV2, err = NewTaskStatsServiceV2()
+	if err != nil {
+		return nil, err
+	}
+	return _serviceV2, nil
+}