318 lines
12 KiB
Go
318 lines
12 KiB
Go
// Package storage provides the persistent index and site-info storage backed by bbolt.
|
||
// storage 包提供基于 bbolt 的持久化存储，负责保存倒排索引、URL摘要缓存和网站元信息。
|
||
//
|
||
// 索引空间（index bucket）：key = 关键词（string），value = brotli 压缩的 JSON 数组，每项为 [权重, URL] 对。
|
||
// 融合之门（gate bucket）：key = URL（string），value = brotli 压缩的 JSON 数组 [标题, 描述, 正文, 时间戳]。
|
||
// 网站之门（site_gate bucket）：key = 主机名（string），value = brotli 压缩的 JSON SiteInfo 结构体。
|
||
//
|
||
// Python 版使用自定义哈希桶结构；Go 版直接交由 bbolt 原生处理。
|
||
package storage
|
||
|
||
import (
|
||
"encoding/json" // JSON 序列化/反序列化
|
||
"fmt" // 格式化错误信息
|
||
"io" // io.EOF 常量
|
||
"os" // 操作系统功能（创建目录等）
|
||
"path/filepath" // 路径拼接
|
||
|
||
"github.com/andybalholm/brotli" // Brotli 无损压缩库（用于压缩存储数据）
|
||
bolt "go.etcd.io/bbolt" // BoltDB，纯 Go 嵌入式 KV 数据库
|
||
)
|
||
|
||
// IndexEntry 是倒排索引中的单个条目。
|
||
// 一条索引记录表示"某个 URL 与某个关键词的相关性权重"。
|
||
type IndexEntry struct {
|
||
Weight float32 `json:"w"` // 该 URL 在该关键词下的得分/权重
|
||
URL string `json:"u"` // 网页 URL
|
||
}
|
||
|
||
// SnippetEntry 是 URL 对应的摘要信息缓存。
|
||
// 包含页面标题、描述、正文片段和抓取时间戳。
|
||
type SnippetEntry struct {
|
||
Title string `json:"title"` // 网页标题
|
||
Description string `json:"desc"` // meta description 或自动生成的描述
|
||
Text string `json:"text"` // 正文前 N 字符的文本片段
|
||
Timestamp int64 `json:"ts"` // 抓取该页面时的 Unix 时间戳
|
||
}
|
||
|
||
// 三个 bbolt bucket 的名称（以字节数组存储，bbolt 要求 key/value 均为字节）
|
||
var (
|
||
bucketIndex = []byte("index") // 倒排索引 bucket
|
||
bucketGate = []byte("gate") // URL 摘要缓存 bucket
|
||
bucketSiteGate = []byte("site_gate") // 网站元信息 bucket
|
||
)
|
||
|
||
// DB 封装一个 bbolt 数据库，提供类型化的存取接口。
|
||
// bbolt 内部已实现并发安全，无需额外加锁。
|
||
type DB struct {
|
||
db *bolt.DB // 底层 bbolt 数据库句柄
|
||
}
|
||
|
||
// Open 在指定目录路径下创建或打开 bbolt 数据库文件。
|
||
// 如果目录不存在会自动创建。数据库文件名为 sese.db。
|
||
func Open(dir string) (*DB, error) {
|
||
// 确保存储目录存在（0775 权限：所有者读写执行，组用户读执行，其他读执行）
|
||
if err := os.MkdirAll(dir, 0o755); err != nil {
|
||
return nil, fmt.Errorf("storage.Open mkdir: %w", err)
|
||
}
|
||
// 拼接数据库文件路径：dir/sese.db
|
||
path := filepath.Join(dir, "sese.db")
|
||
// 打开/创建数据库文件，文件权限 0600（仅所有者可读写）
|
||
db, err := bolt.Open(path, 0o600, nil)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("storage.Open bolt: %w", err)
|
||
}
|
||
// 启动时确保三个 bucket 都存在（不存在则创建）
|
||
err = db.Update(func(tx *bolt.Tx) error {
|
||
for _, b := range [][]byte{bucketIndex, bucketGate, bucketSiteGate} {
|
||
if _, err := tx.CreateBucketIfNotExists(b); err != nil {
|
||
return err
|
||
}
|
||
}
|
||
return nil
|
||
})
|
||
if err != nil {
|
||
return nil, fmt.Errorf("storage.Open create buckets: %w", err)
|
||
}
|
||
return &DB{db: db}, nil
|
||
}
|
||
|
||
// Close 关闭底层 bbolt 数据库连接。
|
||
func (d *DB) Close() error {
|
||
return d.db.Close()
|
||
}
|
||
|
||
// ---- 辅助函数：压缩与解压 ----
|
||
|
||
// compress 将字节数组用 brotli 压缩后返回。
|
||
// brotli 压缩比高于 gzip，适合大量文本的存储空间优化。
|
||
func compress(data []byte) ([]byte, error) {
|
||
buf := make([]byte, 0, len(data)) // 预分配，避免反复扩容
|
||
w := brotli.NewWriterLevel((*appendWriter)(&buf), 6) // 压缩级别 6（平衡速度和压缩比）
|
||
if _, err := w.Write(data); err != nil {
|
||
return nil, err
|
||
}
|
||
if err := w.Close(); err != nil {
|
||
return nil, err
|
||
}
|
||
return buf, nil
|
||
}
|
||
|
||
// decompress 将 brotli 压缩的字节数组解压还原。
|
||
func decompress(data []byte) ([]byte, error) {
|
||
// brotli.NewReader 从字节数组创建读取器（通过 byteReader 适配 io.Reader 接口）
|
||
r := brotli.NewReader(
|
||
(*byteReader)(&data),
|
||
)
|
||
out := make([]byte, 0, len(data)*3) // 预分配约 3 倍空间（解压后通常更大）
|
||
tmp := make([]byte, 4096) // 每次最多读 4KB
|
||
for {
|
||
n, err := r.Read(tmp)
|
||
out = append(out, tmp[:n]...) // 追加本次读取的字节
|
||
if err != nil {
|
||
if err == io.EOF {
|
||
break // 正常读完
|
||
}
|
||
return out, err // 其他错误（非 EOF）则返回
|
||
}
|
||
}
|
||
return out, nil
|
||
}
|
||
|
||
// appendWriter 将 *[]byte 适配为 io.Writer 接口（写入时直接 append）。
|
||
type appendWriter []byte
|
||
|
||
// Write 将数据 p 追加到 appendWriter 末尾，返回写入字节数。
|
||
func (a *appendWriter) Write(p []byte) (int, error) {
|
||
*a = append(*a, p...)
|
||
return len(p), nil
|
||
}
|
||
|
||
// byteReader 将 []byte 适配为 io.Reader 接口（顺序读取，支持读完后返回 EOF）。
|
||
type byteReader []byte
|
||
|
||
// Read 从字节数组读取最多 len(p) 字节到 p 中，返回实际读取字节数和可能的错误。
|
||
// 当字节数组全部读完后返回 io.EOF。
|
||
func (b *byteReader) Read(p []byte) (int, error) {
|
||
if len(*b) == 0 {
|
||
return 0, io.EOF // 已读完
|
||
}
|
||
n := copy(p, *b) // 复制最多 len(p) 字节
|
||
*b = (*b)[n:] // 前进指针
|
||
return n, nil
|
||
}
|
||
|
||
// marshalCompress 将任意可序列化对象先 JSON 编码，再 brotli 压缩，返回压缩后的字节。
|
||
func marshalCompress(v any) ([]byte, error) {
|
||
raw, err := json.Marshal(v) // 先序列化为 JSON
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
return compress(raw) // 再压缩
|
||
}
|
||
|
||
// decompressUnmarshal 将压缩字节先解压，再 JSON 反序列化到目标对象 v。
|
||
func decompressUnmarshal(data []byte, v any) error {
|
||
raw, err := decompress(data) // 先解压
|
||
if err != nil {
|
||
return err
|
||
}
|
||
return json.Unmarshal(raw, v) // 再反序列化
|
||
}
|
||
|
||
// ---- 倒排索引（Index）相关方法 ----
|
||
|
||
// GetIndex 根据关键词查询倒排索引，返回该词关联的所有 [权重, URL] 条目列表。
|
||
func (d *DB) GetIndex(keyword string) ([]IndexEntry, error) {
|
||
var entries []IndexEntry
|
||
err := d.db.View(func(tx *bolt.Tx) error {
|
||
b := tx.Bucket(bucketIndex)
|
||
v := b.Get([]byte(keyword)) // 在 index bucket 中按关键词查询
|
||
if v == nil {
|
||
return nil // 不存在该词，返回空列表
|
||
}
|
||
return decompressUnmarshal(v, &entries)
|
||
})
|
||
return entries, err
|
||
}
|
||
|
||
// SetIndex 将某关键词的完整索引条目列表覆盖写入（替换旧数据）。
|
||
func (d *DB) SetIndex(keyword string, entries []IndexEntry) error {
|
||
data, err := marshalCompress(entries)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
return d.db.Update(func(tx *bolt.Tx) error {
|
||
return tx.Bucket(bucketIndex).Put([]byte(keyword), data)
|
||
})
|
||
}
|
||
|
||
// BatchSetIndex 在一次事务中批量写入多个关键词→条目列表的映射。
|
||
// 比多次调用 SetIndex 效率更高（减少事务开销）。
|
||
func (d *DB) BatchSetIndex(batch map[string][]IndexEntry) error {
|
||
return d.db.Update(func(tx *bolt.Tx) error {
|
||
b := tx.Bucket(bucketIndex)
|
||
for keyword, entries := range batch {
|
||
data, err := marshalCompress(entries)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
if err := b.Put([]byte(keyword), data); err != nil {
|
||
return err
|
||
}
|
||
}
|
||
return nil
|
||
})
|
||
}
|
||
|
||
// ForEachIndex 遍历倒排索引中所有关键词及其关联条目，对每个条目调用 fn。
|
||
// 用于全量读取索引、做备份或重新计算等场景。
|
||
func (d *DB) ForEachIndex(fn func(keyword string, entries []IndexEntry) error) error {
|
||
return d.db.View(func(tx *bolt.Tx) error {
|
||
return tx.Bucket(bucketIndex).ForEach(func(k, v []byte) error {
|
||
var entries []IndexEntry
|
||
if err := decompressUnmarshal(v, &entries); err != nil {
|
||
return nil // 跳过损坏条目，不中断遍历
|
||
}
|
||
return fn(string(k), entries)
|
||
})
|
||
})
|
||
}
|
||
|
||
// ---- 融合之门（Gate）：URL 摘要缓存相关方法 ----
|
||
|
||
// GetSnippet 根据 URL 查询缓存的摘要信息（标题/描述/正文片段）。
|
||
// 若未命中返回 error。
|
||
func (d *DB) GetSnippet(url string) (*SnippetEntry, error) {
|
||
var entry SnippetEntry
|
||
err := d.db.View(func(tx *bolt.Tx) error {
|
||
v := tx.Bucket(bucketGate).Get([]byte(url))
|
||
if v == nil {
|
||
return fmt.Errorf("not found")
|
||
}
|
||
return decompressUnmarshal(v, &entry)
|
||
})
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
return &entry, nil
|
||
}
|
||
|
||
// SetSnippet 将某 URL 的摘要信息写入缓存（覆盖已有数据）。
|
||
func (d *DB) SetSnippet(url string, entry *SnippetEntry) error {
|
||
data, err := marshalCompress(entry)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
return d.db.Update(func(tx *bolt.Tx) error {
|
||
return tx.Bucket(bucketGate).Put([]byte(url), data)
|
||
})
|
||
}
|
||
|
||
// ---- 网站之门（SiteGate）：网站元信息相关方法 ----
|
||
|
||
// SiteInfo 存放每个域名/主机的元信息，与 Python 版网站.py 的 dataclass 对应。
|
||
type SiteInfo struct {
|
||
VisitCount int `json:"visit_count"` // 累计访问该网站的次数
|
||
LastVisitTime int64 `json:"last_visit_time"` // 上次访问该网站的时间戳
|
||
Fingerprint any `json:"fingerprint,omitempty"` // 网站指纹（用于识别重复站点）
|
||
SuccessRate *float64 `json:"success_rate,omitempty"` // 访问成功率（成功次数/总访问次数）
|
||
HTMLStructure string `json:"html_structure,omitempty"` // HTML 结构特征摘要
|
||
IPs []string `json:"ips,omitempty"` // 该域名解析出的 IP 列表
|
||
Quality *float64 `json:"quality,omitempty"` // 网站质量评分（0~1）
|
||
HTTPSAvailable *bool `json:"https_available,omitempty"` // 是否支持 HTTPS
|
||
Keywords []string `json:"keywords,omitempty"` // 该网站的高频关键词列表
|
||
OutLinks []string `json:"out_links,omitempty"` // 从该网站页面提取的出站链接列表
|
||
Languages map[string]float64 `json:"languages,omitempty"` // 网站语种分布（语种代码 → 占比）
|
||
Redirects map[string]string `json:"redirects,omitempty"` // 重定向链（URL → 最终 URL）
|
||
ServerTypes []string `json:"server_types,omitempty"` // 网站使用的 HTTP Server 类型列表
|
||
}
|
||
|
||
// GetSiteInfo 根据主机名查询网站元信息。
|
||
// 若不存在则返回仅有默认空 map 的空 SiteInfo（不报错，方便调用方直接使用）。
|
||
func (d *DB) GetSiteInfo(host string) (*SiteInfo, error) {
|
||
var info SiteInfo
|
||
err := d.db.View(func(tx *bolt.Tx) error {
|
||
v := tx.Bucket(bucketSiteGate).Get([]byte(host))
|
||
if v == nil {
|
||
return fmt.Errorf("not found")
|
||
}
|
||
return decompressUnmarshal(v, &info)
|
||
})
|
||
if err != nil {
|
||
// 未找到时返回带默认空 map 的空结构，避免调用方空指针 panic
|
||
return &SiteInfo{Languages: make(map[string]float64), Redirects: make(map[string]string)}, nil
|
||
}
|
||
if info.Languages == nil {
|
||
info.Languages = make(map[string]float64)
|
||
}
|
||
if info.Redirects == nil {
|
||
info.Redirects = make(map[string]string)
|
||
}
|
||
return &info, nil
|
||
}
|
||
|
||
// SetSiteInfo 将某主机名的网站元信息写入存储（覆盖已有数据）。
|
||
func (d *DB) SetSiteInfo(host string, info *SiteInfo) error {
|
||
data, err := marshalCompress(info)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
return d.db.Update(func(tx *bolt.Tx) error {
|
||
return tx.Bucket(bucketSiteGate).Put([]byte(host), data)
|
||
})
|
||
}
|
||
|
||
// ForEachSite 遍历所有网站元信息条目，对每个条目调用 fn。
|
||
func (d *DB) ForEachSite(fn func(host string, info *SiteInfo) error) error {
|
||
return d.db.View(func(tx *bolt.Tx) error {
|
||
return tx.Bucket(bucketSiteGate).ForEach(func(k, v []byte) error {
|
||
var info SiteInfo
|
||
if err := decompressUnmarshal(v, &info); err != nil {
|
||
return nil // 跳过损坏条目
|
||
}
|
||
return fn(string(k), &info)
|
||
})
|
||
})
|
||
}
|