Files
sese-engine-go/config/config.go
T
2026-04-12 14:24:26 +08:00

360 lines
12 KiB
Go

// Package config holds all global configuration parameters for sese-engine.
// config 包存放 sese-engine 的所有全局配置参数。
package config
import (
"fmt"
"math"
"os"
"path/filepath"
"reflect"
"sync"
"gopkg.in/yaml.v3"
)
// configMu 保护 Global 的运行时修改(动态调参场景)。
var configMu sync.RWMutex
// Config 是完整的配置结构体
type Config struct {
Index IndexConfig `yaml:"index"`
Crawler CrawlerConfig `yaml:"crawler"`
Search SearchConfig `yaml:"search"`
Backlink BacklinkConfig `yaml:"backlink"`
Storage StorageConfig `yaml:"storage"`
Prometheus PrometheusConfig `yaml:"prometheus"`
}
// IndexConfig 索引/存储相关限制
type IndexConfig struct {
MaxURLsPerKey int `yaml:"max_urls_per_key"`
MaxSameDomainPerKey int `yaml:"max_same_domain_per_key"`
BigCleanThreshold int `yaml:"big_clean_threshold"`
MaxNewURLsPerKey int `yaml:"max_new_urls_per_key"`
MinURLsForNewKey int `yaml:"min_urls_for_new_key"`
}
// CrawlerConfig 爬虫行为相关配置
type CrawlerConfig struct {
SpiderName string `yaml:"spider_name"`
Cooldown int `yaml:"cooldown"`
Workers int `yaml:"workers"`
CrawlFocus float64 `yaml:"crawl_focus"`
MaxKeywordsPerPage int `yaml:"max_keywords_per_page"`
MaxEpoch int `yaml:"max_epoch"`
ExpectedProsperRatio float64 `yaml:"expected_prosper_ratio"`
EntryURL string `yaml:"entry_url"`
MaxPageSize int `yaml:"max_page_size"` // 单个页面最大抓取字节数(0=不限,默认 5MB)
RecrawlMaxAge int `yaml:"recrawl_max_age"` // URL 过期时间(秒),超过此时间的 URL 允许被重爬,默认 30 天
RecrawlCheckInterval int `yaml:"recrawl_check_interval"` // 运行期间检查过期 URL 的间隔(秒),默认 1 小时
RecrawlBatchSize int `yaml:"recrawl_batch_size"` // 每次检查最多释放多少个过期 URL,默认 500
MaxPriorityChildren int `yaml:"max_priority_children"` // priorityChildren 队列的最大链接数,默认 100
}
// SearchConfig 搜索结果排序权重配置
type SearchConfig struct {
UseOnlineSnippet bool `yaml:"use_online_snippet"`
OnlineSnippetTimeout int `yaml:"online_snippet_timeout"`
WeightDailyDecay float64 `yaml:"weight_daily_decay"`
LanguageWeight float64 `yaml:"language_weight"`
ConsecutiveKeyWeight float64 `yaml:"consecutive_key_weight"`
BacklinkWeight float64 `yaml:"backlink_weight"`
ServerPort int `yaml:"server_port"`
FlushIntervalSeconds int `yaml:"flush_interval_seconds"`
StatsRefreshInterval int `yaml:"stats_refresh_interval"` // 统计缓存刷新间隔(秒),默认 30
MissPenalty float64 `yaml:"miss_penalty"` // 缺词惩罚系数(0=不惩罚,1=完全忽略缺词URL),默认 0.15
UnixSocket string `yaml:"unix_socket"` // Unix socket 路径(仅 Linux/macOS),空字符串表示不启用
}
// BacklinkConfig 反向链接计算相关配置
type BacklinkConfig struct {
Baseline int `yaml:"baseline"`
}
// StorageConfig 存储配置
type StorageConfig struct {
Path string `yaml:"path"`
}
// PrometheusConfig Prometheus监控端口配置
type PrometheusConfig struct {
CrawlerPort int `yaml:"crawler_port"`
BacklinkPort int `yaml:"backlink_port"`
SearchPort int `yaml:"search_port"`
}
// Global 全局配置实例,加载后可通过此变量访问
var Global Config
// Load 从指定路径加载配置文件,并自动补全缺失的字段。
// 流程:读取 YAML → 与默认值合并 → 写回 config.yml → 赋值 Global
func Load(configPath string) error {
data, err := os.ReadFile(configPath)
if err != nil {
return fmt.Errorf("failed to read config file: %v", err)
}
// 先拿到当前 YAML 内容,用于判断哪些字段实际存在于文件中
var yamlOnly Config
if err := yaml.Unmarshal(data, &yamlOnly); err != nil {
return fmt.Errorf("failed to parse config file: %v", err)
}
// 从默认值开始,YAML 中有值的字段会被覆盖
merged := GetDefaultConfig()
mergeConfig(&merged, &yamlOnly)
// 写回 config.yml(自动补全缺失字段)
yamlOut, err := yaml.Marshal(&merged)
if err != nil {
return fmt.Errorf("failed to marshal config: %v", err)
}
if err := os.WriteFile(configPath, yamlOut, 0644); err != nil {
return fmt.Errorf("failed to write config file: %v", err)
}
Global = merged
return nil
}
// mergeConfig 将 src 中的非零字段合并到 dst(原地修改 dst)。
// 用于把 YAML 实际配置值覆盖到默认值结构上。
func mergeConfig(dst, src interface{}) {
if dst == nil || src == nil {
return
}
dstVal := reflect.ValueOf(dst).Elem()
srcVal := reflect.ValueOf(src).Elem()
for i := 0; i < dstVal.NumField(); i++ {
dstField := dstVal.Field(i)
srcField := srcVal.Field(i)
switch dstField.Kind() {
case reflect.Struct:
// 递归合并嵌套 struct
mergeConfig(dstField.Addr().Interface(), srcField.Addr().Interface())
case reflect.Slice:
// slice:仅当 src 非空时才覆盖(避免覆盖用户显式设置的长 0 slice)
if srcField.Len() > 0 {
dstField.Set(srcField)
}
default:
// 其他类型:src 为零值则保留 dst 原值(默认值)
if !isZero(srcField) {
dstField.Set(srcField)
}
}
}
}
// isZero 检查 reflect.Value 是否为该类型的零值。
func isZero(v reflect.Value) bool {
switch v.Kind() {
case reflect.Bool:
return !v.Bool()
case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
return v.Int() == 0
case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
return v.Uint() == 0
case reflect.Float32, reflect.Float64:
return math.Float64bits(v.Float()) == 0
case reflect.String:
return v.String() == ""
case reflect.Ptr, reflect.Interface:
return v.IsNil()
}
return reflect.DeepEqual(v.Interface(), reflect.Zero(v.Type()).Interface())
}
// LoadFromSavedata 从 savedata 目录加载 config.yml
func LoadFromSavedata() error {
configPath := filepath.Join("savedata", "config.yml")
return Load(configPath)
}
// GetDefaultConfig 返回默认配置
func GetDefaultConfig() Config {
return Config{
Index: IndexConfig{
MaxURLsPerKey: 11000,
MaxSameDomainPerKey: 20,
BigCleanThreshold: 2000000,
MaxNewURLsPerKey: 10000,
MinURLsForNewKey: 3,
},
Crawler: CrawlerConfig{
SpiderName: "Haibara_AI_spider",
Cooldown: 3,
Workers: 22,
CrawlFocus: 0.7,
MaxKeywordsPerPage: 250,
MaxEpoch: 100,
ExpectedProsperRatio: 0.6,
EntryURL: "https://haibara.ai/",
MaxPageSize: 5 * 1024 * 1024,
RecrawlMaxAge: 30 * 86400, // 30 天
RecrawlCheckInterval: 3600, // 1 小时
RecrawlBatchSize: 500,
MaxPriorityChildren: 100,
},
Search: SearchConfig{
UseOnlineSnippet: true,
OnlineSnippetTimeout: 3,
WeightDailyDecay: 0.996,
LanguageWeight: 0.5,
ConsecutiveKeyWeight: 1.3,
BacklinkWeight: 1.0,
ServerPort: 50082,
FlushIntervalSeconds: 300,
StatsRefreshInterval: 30,
MissPenalty: 0.15,
},
Backlink: BacklinkConfig{
Baseline: 200000,
},
Storage: StorageConfig{
Path: "./savedata",
},
Prometheus: PrometheusConfig{
CrawlerPort: 14950,
BacklinkPort: 14952,
SearchPort: 14953,
},
}
}
// 以下是向后兼容的常量定义,使用 Global 变量的值
// 在 Init() 被调用后,这些函数会返回加载的配置值
func init() {
// 初始化时设置默认值
Global = GetDefaultConfig()
}
// MaxURLsPerKey 返回配置值
func MaxURLsPerKey() int { return Global.Index.MaxURLsPerKey }
// MaxSameDomainPerKey 返回配置值
func MaxSameDomainPerKey() int { return Global.Index.MaxSameDomainPerKey }
// BigCleanThreshold 返回配置值
func BigCleanThreshold() int { return Global.Index.BigCleanThreshold }
// MaxNewURLsPerKey 返回配置值
func MaxNewURLsPerKey() int { return Global.Index.MaxNewURLsPerKey }
// MinURLsForNewKey 返回配置值
func MinURLsForNewKey() int { return Global.Index.MinURLsForNewKey }
// SpiderName 返回配置值
func SpiderName() string { return Global.Crawler.SpiderName }
// CrawlerCooldown 返回配置值
func CrawlerCooldown() int { return Global.Crawler.Cooldown }
// CrawlerWorkers 返回配置值
func CrawlerWorkers() int { return Global.Crawler.Workers }
// SetCrawlerWorkers 在运行时动态修改爬虫并发数(线程安全)。
func SetCrawlerWorkers(n int) {
if n < 1 {
n = 1
}
if n > 500 {
n = 500
}
configMu.Lock()
Global.Crawler.Workers = n
configMu.Unlock()
}
// CrawlFocus 返回配置值
func CrawlFocus() float64 { return Global.Crawler.CrawlFocus }
// MaxKeywordsPerPage 返回配置值
func MaxKeywordsPerPage() int { return Global.Crawler.MaxKeywordsPerPage }
// MaxEpoch 返回配置值
func MaxEpoch() int { return Global.Crawler.MaxEpoch }
// ExpectedProsperRatio 返回配置值
func ExpectedProsperRatio() float64 { return Global.Crawler.ExpectedProsperRatio }
// EntryURL 返回配置值
func EntryURL() string { return Global.Crawler.EntryURL }
// MaxPageSize 返回单个页面最大抓取字节数(0=不限)。
func MaxPageSize() int { return Global.Crawler.MaxPageSize }
// RecrawlMaxAge 返回 URL 过期时间(秒),超过此时间的 URL 允许被重爬。
func RecrawlMaxAge() int { return Global.Crawler.RecrawlMaxAge }
// RecrawlCheckInterval 返回运行期间检查过期 URL 的间隔(秒)。
func RecrawlCheckInterval() int { return Global.Crawler.RecrawlCheckInterval }
// RecrawlBatchSize 返回每次检查最多释放的过期 URL 数量。
func RecrawlBatchSize() int { return Global.Crawler.RecrawlBatchSize }
// UseOnlineSnippet 返回配置值
func UseOnlineSnippet() bool { return Global.Search.UseOnlineSnippet }
// OnlineSnippetTimeout 返回配置值
func OnlineSnippetTimeout() int { return Global.Search.OnlineSnippetTimeout }
// WeightDailyDecay 返回配置值
func WeightDailyDecay() float64 { return Global.Search.WeightDailyDecay }
// LanguageWeight 返回配置值
func LanguageWeight() float64 { return Global.Search.LanguageWeight }
// ConsecutiveKeyWeight 返回配置值
func ConsecutiveKeyWeight() float64 { return Global.Search.ConsecutiveKeyWeight }
// BacklinkWeight 返回配置值
func BacklinkWeight() float64 { return Global.Search.BacklinkWeight }
// SearchServerPort 返回配置值
func SearchServerPort() int { return Global.Search.ServerPort }
// FlushIntervalSeconds 返回配置值
func FlushIntervalSeconds() int { return Global.Search.FlushIntervalSeconds }
// StatsRefreshInterval 返回统计缓存刷新间隔(秒),默认 30。
func StatsRefreshInterval() int {
if Global.Search.StatsRefreshInterval <= 0 {
return 30
}
return Global.Search.StatsRefreshInterval
}
// MissPenalty 返回缺词惩罚系数(0~1),值越大对缺少查询词的 URL 惩罚越重。
func MissPenalty() float64 { return Global.Search.MissPenalty }
// UnixSocket 返回 Unix socket 路径,空字符串表示不启用。
func UnixSocket() string { return Global.Search.UnixSocket }
// BacklinkBaseline 返回配置值
func BacklinkBaseline() int { return Global.Backlink.Baseline }
// PromPortCrawler 返回配置值
func PromPortCrawler() int { return Global.Prometheus.CrawlerPort }
// PromPortBacklink 返回配置值
func PromPortBacklink() int { return Global.Prometheus.BacklinkPort }
// PromPortSearch 返回配置值
func PromPortSearch() int { return Global.Prometheus.SearchPort }
// MaxPriorityChildren 返回 priorityChildren 队列的最大链接数(0=不限)。
func MaxPriorityChildren() int {
if Global.Crawler.MaxPriorityChildren <= 0 {
return 100 // 默认 100
}
return Global.Crawler.MaxPriorityChildren
}
// 为了向后兼容,保留 StoragePath 常量
const StoragePath = "./savedata"