如何在Go中动态追踪压缩文件大小并实现文件分割？

用户提供的Go代码片段如下：

func main() {
    // Some text we want to compress.
    original := "bird and frog"
    
    // Open a file for writing.
    f, _ := os.Create("C:\\programs\\file.gz")
    
    // Create gzip writer.
    w := gzip.NewWriter(f)
    
    // Write bytes in compressed form to the file.
    while ( looping over database cursor) {
       w.Write([]byte(/* the row from the database as obtained from cursor */))
    }
    
    // Close the file.
    w.Close()
    
    fmt.Println("DONE")
}

需求：当压缩后的文件大小达到指定阈值时，自动关闭当前文件并新建一个压缩文件继续写入。例如：数据库有10条50字节的数据，压缩比为2（压缩后每条占25字节），文件大小限制为50字节，那么每写入2条数据就需要切换一次文件。需要实现动态追踪压缩后的文件实际大小。

gzip.Writer本身不会直接暴露压缩后写入文件的字节数，因此需要自定义一个统计写入量的包装器，结合文件切换逻辑来实现需求。

修改后的完整代码 #

package main

import (
	"compress/gzip"
	"fmt"
	"io"
	"os"
)

// CountWriter 包装底层Writer，统计实际写入的字节数
type CountWriter struct {
	W     io.Writer
	Count int64
}

// Write 实现io.Writer接口，转发数据的同时累加写入量
func (cw *CountWriter) Write(p []byte) (n int, err error) {
	n, err = cw.W.Write(p)
	cw.Count += int64(n)
	return
}

func main() {
	const threshold = 50 // 压缩后文件的大小阈值（单位：字节）
	fileIndex := 0
	var currentFile *os.File
	var gzWriter *gzip.Writer
	var countWriter *CountWriter

	// 初始化第一个文件
	if err := createNewFile(&currentFile, &gzWriter, &countWriter, &fileIndex); err != nil {
		fmt.Printf("初始化文件失败: %v\n", err)
		return
	}

	// 模拟遍历数据库游标（示例为10条50字节的数据）
	for i := 0; i < 10; i++ {
		rowData := make([]byte, 50) // 模拟每条50字节的原始数据
		if _, err := gzWriter.Write(rowData); err != nil {
			fmt.Printf("写入数据失败: %v\n", err)
			cleanup(currentFile, gzWriter)
			return
		}

		// 检查是否达到阈值，触发文件切换
		if countWriter.Count >= threshold {
			cleanup(currentFile, gzWriter)
			if err := createNewFile(&currentFile, &gzWriter, &countWriter, &fileIndex); err != nil {
				fmt.Printf("创建新文件失败: %v\n", err)
				return
			}
		}
	}

	// 关闭最后一个文件资源
	cleanup(currentFile, gzWriter)
	fmt.Println("DONE")
}

// createNewFile 创建新的压缩文件，初始化gzip写入器和计数器
func createNewFile(f **os.File, gz **gzip.Writer, cw **CountWriter, index *int) error {
	*index++
	fileName := fmt.Sprintf("C:\\programs\\file_%d.gz", *index)

	newFile, err := os.Create(fileName)
	if err != nil {
		return err
	}

	newCountWriter := &CountWriter{W: newFile, Count: 0}
	newGzWriter := gzip.NewWriter(newCountWriter)

	*f = newFile
	*gz = newGzWriter
	*cw = newCountWriter
	return nil
}

// cleanup 安全关闭文件和gzip写入器，避免资源泄漏和压缩文件损坏
func cleanup(f *os.File, gz *gzip.Writer) {
	if gz != nil {
		_ = gz.Close()
	}
	if f != nil {
		_ = f.Close()
	}
}

关键逻辑说明 #

1. CountWriter统计压缩后字节数：在gzip.Writer和底层文件之间插入该包装器，实时统计实际写入文件的压缩后数据量，确保判断阈值的准确性。
2. 文件切换触发：每次写入数据后检查计数器，达到阈值时先关闭当前的gzip.Writer和文件（必须先关gzip，否则压缩文件会损坏），再调用createNewFile初始化新的写入资源。
3. 资源安全管理：通过cleanup函数统一处理资源关闭，避免出现文件句柄泄漏或压缩文件不完整的问题。

注意事项 #

• 压缩比是动态变化的，不能用原始数据大小预估压缩后体积，必须依赖实际写入量统计。
• 每次切换文件时必须先关闭gzip.Writer，它会自动写入gzip格式的尾部校验信息，保证文件可正常解压。

内容的提问来源于stack exchange，提问作者JavaDeveloper