ragflow/internal/syncer/syncer.go

//
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
//
// 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 syncer

import (
	"context"
	"errors"
	"fmt"
	"sync"
	"time"

	"ragflow/internal/common"
	"ragflow/internal/dao"
	"ragflow/internal/engine"
	"ragflow/internal/service"
	documentservice "ragflow/internal/service/document"
	syncerconnector "ragflow/internal/syncer/connector"
	"ragflow/internal/utility"

	"go.uber.org/zap"
)

// Syncer owns NATS/DB scheduling, task workers, and the shared batch job executor.
type Syncer struct {
	id          string
	config      Config
	queue       chan TaskEnvelope
	scheduler   *Scheduler
	worker      *TaskWorker
	executor    *SyncJobExecutor
	cancel      context.CancelFunc
	workerGroup sync.WaitGroup
	stopOnce    sync.Once
	ShutdownCh  chan struct{}
}

// NewSyncer creates a server-compatible syncer with default dependencies.
func NewSyncer(taskWorkerCount int, pollInterval time.Duration) *Syncer {
	// init the config
	config := DefaultConfig()
	config.TaskWorkerCount = taskWorkerCount
	config.PollInterval = pollInterval

	taskDAO := dao.NewSyncTaskDAO(nil)
	registry := syncerconnector.NewRegistry()

	registerBuiltInConnectors(registry)

	documentService := documentservice.NewDocumentService()
	pruneService := service.NewSyncPruneService(documentService, nil)

	return New(config, taskDAO, registry, documentService, pruneService)
}

// New creates a datasource syncer from explicit dependencies.
func New(config Config, taskDAO *dao.SyncTaskDAO, registry ConnectorRegistry, sink service.DocumentSink, pruneService *service.SyncPruneService) *Syncer {
	config = config.Normalize()
	queue := make(chan TaskEnvelope, config.TaskQueueSize)
	locker := NewConnectorLock()

	executor := NewSyncJobExecutor(SyncJobExecutorConfig{
		WorkerCount:  config.JobWorkerCount,
		JobQueueSize: config.JobQueueSize,
	})
	taskService := service.NewSyncTaskService(taskDAO)
	idResolver := service.NewDocumentIDResolver(service.NewGormDocumentStore())

	coordinator := NewTaskCoordinator(TaskCoordinatorConfig{
		ItemRetryCount:     config.ItemRetryCount,
		ItemRetryBaseDelay: config.ItemRetryBaseDelay,
	}, taskService, registry, sink, pruneService, idResolver, executor)

	scheduler := NewScheduler(config.PollInterval, queue, taskService)
	if broker, ok := engine.GetMessageQueueEngine().(SyncTaskBroker); ok {
		scheduler = NewNATSScheduler(config.PollInterval, queue, taskService, broker)
	}

	return &Syncer{
		id:         utility.GenerateUUID(),
		config:     config,
		queue:      queue,
		scheduler:  scheduler,
		worker:     NewTaskWorker(queue, taskService, coordinator, locker),
		executor:   executor,
		ShutdownCh: make(chan struct{}),
	}
}

// ID returns this syncer process ID.
func (s *Syncer) ID() string {
	if s == nil {
		return ""
	}
	return s.id
}

// Start launches the scheduler and task workers with a background context.
func (s *Syncer) Start() error {
	return s.StartContext(context.Background())
}

// StartContext launches the scheduler and task workers.
func (s *Syncer) StartContext(ctx context.Context) error {
	if s == nil {
		return errors.New("syncer is nil")
	}

	runCtx, cancel := context.WithCancel(ctx)
	s.cancel = cancel
	s.workerGroup.Add(2)

	// run scheduler
	go func() {
		defer s.workerGroup.Done()
		if err := s.scheduler.Run(runCtx); err != nil && !errors.Is(err, context.Canceled) {
			common.Error("syncer scheduler stopped", err)
		}
	}()

	// run worker poll
	go func() {
		defer s.workerGroup.Done()
		s.worker.Run(runCtx, s.config.TaskWorkerCount)
	}()
	return nil
}

// Stop cancels the scheduler and waits for workers to exit.
func (s *Syncer) Stop() {
	if s == nil {
		return
	}
	s.stopOnce.Do(func() {
		if s.cancel != nil {
			s.cancel()
		}
		s.workerGroup.Wait()
		s.executor.Close()
		close(s.ShutdownCh)
	})
}

// logSyncTaskDuration test run time, delete it soon
func logSyncTaskDuration(taskContext service.SyncTaskContext, startedAt time.Time) {
	if taskContext.Task.TaskType != service.TaskTypeSync {
		return
	}
	common.Info(
		"sync task duration",
		zap.String("task_id", taskContext.Task.ID),
		zap.String("connector_id", taskContext.Connector.ID),
		zap.String("kb_id", taskContext.Knowledgebase.ID),
		zap.String("source", taskContext.Connector.Source),
		zap.Duration("elapsed", time.Since(startedAt)),
	)
}

// registerBuiltInConnectors registers datasource connectors available in the server binary.
func registerBuiltInConnectors(registry *syncerconnector.Registry) {
	registerDAOConnector(registry, "rss", syncerconnector.NewRSSConnector)
	registerDAOConnector(registry, "github", syncerconnector.NewGitHubConnector)
	registerDAOConnector(registry, "gmail", syncerconnector.NewGmailConnector)
	registerDAOConnector(registry, "google-drive", syncerconnector.NewGoogleDriveConnector)
	registerDAOConnector(registry, "google_drive", syncerconnector.NewGoogleDriveConnector)
}

func registerDAOConnector[T syncerconnector.Connector](registry *syncerconnector.Registry, source string, factory func(map[string]any) (T, error)) {
	registry.Register(source, func(ctx context.Context, taskContext any) (syncerconnector.Connector, error) {
		row, ok := taskContext.(dao.SyncTaskContext)
		if !ok {
			return nil, fmt.Errorf("%s connector received an invalid task context", source)
		}
		return factory(map[string]any(row.Connector.Config))
	})
}