Files
sese-engine-go/config/config.go
T
2026-04-20 18:26:54 +08:00

553 lines
17 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"`
MySQL MySQLConfig `yaml:"mysql"`
Redis RedisConfig `yaml:"redis"`
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"`
}
// MySQLConfig MySQL数据库连接配置
// 支持两种连接方式:Unix Socket 和 TCP
// 优先级:UnixSocket > TCP(如果UnixSocket非空则优先使用)
type MySQLConfig struct {
// 是否启用 MySQL(默认关闭)
Enabled bool `yaml:"enabled"`
// 连接方式: "socket" 或 "tcp"(自动推断,可不填)
Network string `yaml:"network"`
// Unix Socket 路径(Linux/macOS),优先使用
// 示例: "/var/run/mysqld/mysqld.sock" 或 "/tmp/mysql.sock"
UnixSocket string `yaml:"unix_socket"`
// TCP 连接方式:服务器地址
Host string `yaml:"host"`
// TCP 连接方式:端口号
Port int `yaml:"port"`
// 用户名
User string `yaml:"user"`
// 密码
Password string `yaml:"password"`
// 数据库名
Database string `yaml:"database"`
// 连接超时时间(秒)
ConnMaxLifetime int `yaml:"conn_max_lifetime"` // 秒
// 最大空闲连接数
MaxIdleConns int `yaml:"max_idle_conns"`
// 最大打开连接数
MaxOpenConns int `yaml:"max_open_conns"`
}
// RedisConfig Redis连接配置
// 支持两种连接方式:Unix Socket 和 TCP
// 优先级:UnixSocket > TCP(如果UnixSocket非空则优先使用)
type RedisConfig struct {
// 连接方式: "socket" 或 "tcp"(自动推断,可不填)
Network string `yaml:"network"`
// Unix Socket 路径,优先使用
// 示例: "/var/run/redis/redis.sock" 或 "/tmp/redis.sock"
UnixSocket string `yaml:"unix_socket"`
// TCP 连接方式:服务器地址
Host string `yaml:"host"`
// TCP 连接方式:端口号
Port int `yaml:"port"`
// 密码(无密码则留空)
Password string `yaml:"password"`
// 数据库编号(0-15),默认 15
DB int `yaml:"db"`
// 池大小(最大连接数)
PoolSize int `yaml:"pool_size"`
// 最小空闲连接数
MinIdleConns int `yaml:"min_idle_conns"`
// 读超时时间(毫秒)
ReadTimeout int `yaml:"read_timeout"` // 毫秒
// 写超时时间(毫秒)
WriteTimeout int `yaml:"write_timeout"` // 毫秒
}
// 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",
},
MySQL: MySQLConfig{
Enabled: false,
Network: "tcp",
UnixSocket: "",
Host: "localhost",
Port: 3306,
User: "root",
Password: "",
Database: "sese_engine",
ConnMaxLifetime: 3600, // 1小时
MaxIdleConns: 10,
MaxOpenConns: 100,
},
Redis: RedisConfig{
Network: "tcp",
UnixSocket: "",
Host: "localhost",
Port: 6379,
Password: "",
DB: 15, // 默认使用15号数据库
PoolSize: 100,
MinIdleConns: 10,
ReadTimeout: 500, // 毫秒
WriteTimeout: 500, // 毫秒
},
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"
// ---- MySQL 配置访问函数 ----
// MySQLEnabled 返回是否启用 MySQL(默认关闭)
func MySQLEnabled() bool {
return Global.MySQL.Enabled
}
// MySQLDSN 返回 MySQL 连接字符串(DSN)
// 根据配置自动选择 Unix Socket 或 TCP 方式
func MySQLDSN() string {
cfg := Global.MySQL
if cfg.UnixSocket != "" {
// 使用 Unix Socket 连接(推荐,本地连接性能更好)
return cfg.User + ":" + cfg.Password + "@unix(" + cfg.UnixSocket + ")/" + cfg.Database + "?parseTime=true&loc=Local"
}
// 使用 TCP 连接
return cfg.User + ":" + cfg.Password + "@tcp(" + cfg.Host + ":" + itoa(cfg.Port) + ")/" + cfg.Database + "?parseTime=true&loc=Local"
}
// MySQLConnMaxLifetime 返回连接最大生命周期(秒)
func MySQLConnMaxLifetime() int {
if Global.MySQL.ConnMaxLifetime <= 0 {
return 3600
}
return Global.MySQL.ConnMaxLifetime
}
// MySQLMaxIdleConns 返回最大空闲连接数
func MySQLMaxIdleConns() int {
if Global.MySQL.MaxIdleConns <= 0 {
return 10
}
return Global.MySQL.MaxIdleConns
}
// MySQLMaxOpenConns 返回最大打开连接数
func MySQLMaxOpenConns() int {
if Global.MySQL.MaxOpenConns <= 0 {
return 100
}
return Global.MySQL.MaxOpenConns
}
// ---- Redis 配置访问函数 ----
// RedisAddr 返回 Redis 连接地址
// 根据配置自动选择 Unix Socket 或 TCP 方式
func RedisAddr() string {
cfg := Global.Redis
if cfg.UnixSocket != "" {
return cfg.UnixSocket
}
return cfg.Host + ":" + itoa(cfg.Port)
}
// RedisPoolSize 返回连接池大小
func RedisPoolSize() int {
if Global.Redis.PoolSize <= 0 {
return 100
}
return Global.Redis.PoolSize
}
// RedisMinIdleConns 返回最小空闲连接数
func RedisMinIdleConns() int {
if Global.Redis.MinIdleConns <= 0 {
return 10
}
return Global.Redis.MinIdleConns
}
// RedisDB 返回数据库编号
func RedisDB() int {
return Global.Redis.DB
}
// RedisPassword 返回密码(空字符串表示无密码)
func RedisPassword() string {
return Global.Redis.Password
}
// RedisReadTimeout 返回读超时时间(毫秒)
func RedisReadTimeout() int {
if Global.Redis.ReadTimeout <= 0 {
return 500
}
return Global.Redis.ReadTimeout
}
// RedisWriteTimeout 返回写超时时间(毫秒)
func RedisWriteTimeout() int {
if Global.Redis.WriteTimeout <= 0 {
return 500
}
return Global.Redis.WriteTimeout
}
// itoa 将 int 转换为字符串(避免导入 strconv)
func itoa(n int) string {
if n == 0 {
return "0"
}
result := ""
for n > 0 {
result = string(rune('0'+n%10)) + result
n /= 10
}
return result
}