up

storage/storage.go

@@ -410,8 +410,28 @@ type SiteInfo struct {
 }
 
 // GetSiteInfo 根据主机名查询网站元信息。
+// 优先从写缓冲中读取（保证与最近 SetSiteInfo 的数据一致），未命中再读 bbolt。
 // 若不存在则返回仅有默认空 map 的空 SiteInfo（不报错，方便调用方直接使用）。
 func (d *DB) GetSiteInfo(host string) (*SiteInfo, error) {
+	// 先检查写缓冲中是否有该 host 的最新待写入数据
+	d.writeBufMu.Lock()
+	if op, ok := d.writeBuf["site:"+host]; ok {
+		d.writeBufMu.Unlock()
+		var info SiteInfo
+		if err := decompressUnmarshal(op.data, &info); err == nil {
+			if info.Languages == nil {
+				info.Languages = make(map[string]float64)
+			}
+			if info.Redirects == nil {
+				info.Redirects = make(map[string]string)
+			}
+			return &info, nil
+		}
+		// 反序列化失败（不应发生），fall through 到读 db
+	} else {
+		d.writeBufMu.Unlock()
+	}
+
 	var info SiteInfo
 	err := d.db.View(func(tx *bolt.Tx) error {
 		v := tx.Bucket(bucketSiteGate).Get([]byte(host))
@@ -450,6 +470,62 @@ func (d *DB) SetSiteInfo(host string, info *SiteInfo) error {
 	return nil
 }
 
+// UpdateSiteInfo 原子地读取当前 SiteInfo 并应用修改函数 fn，然后写回。
+// 整个读-改-写过程在 writeBufMu 锁内完成，消除并发 lost update 竞态。
+// 适用于多个 goroutine 对同一 host 的 SiteInfo 进行读-改-写的场景。
+func (d *DB) UpdateSiteInfo(host string, fn func(*SiteInfo)) error {
+	d.writeBufMu.Lock()
+
+	// 从缓冲或 db 读取最新 SiteInfo
+	var info *SiteInfo
+	if op, ok := d.writeBuf["site:"+host]; ok {
+		var si SiteInfo
+		if err := decompressUnmarshal(op.data, &si); err == nil {
+			info = &si
+		}
+	}
+
+	if info == nil {
+		var si 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, &si)
+		})
+		if err != nil {
+			info = &SiteInfo{Languages: make(map[string]float64), Redirects: make(map[string]string)}
+		} else {
+			info = &si
+		}
+	}
+	if info.Languages == nil {
+		info.Languages = make(map[string]float64)
+	}
+	if info.Redirects == nil {
+		info.Redirects = make(map[string]string)
+	}
+
+	// 在锁内调用修改函数
+	fn(info)
+
+	// 序列化并写回缓冲
+	data, err := marshalCompress(info)
+	if err != nil {
+		d.writeBufMu.Unlock()
+		return err
+	}
+	d.writeBuf["site:"+host] = &writeOp{opType: 1, key: host, data: data}
+	if len(d.writeBuf) >= 5000 {
+		d.writeBufMu.Unlock()
+		d.flushWriteBuf()
+		return nil
+	}
+	d.writeBufMu.Unlock()
+	return nil
+}
+
 // ForEachSite 遍历所有网站元信息条目，对每个条目调用 fn。
 func (d *DB) ForEachSite(fn func(host string, info *SiteInfo) error) error {
 	return d.db.View(func(tx *bolt.Tx) error {