Go语言中高效跳过io.Reader流中指定字节数的方法

package main

import (
    "fmt"
    "io"
    "strings"
)

// SkipNBytesFromReader 从给定的io.Reader中跳过指定数量的字节
func SkipNBytesFromReader(r io.Reader, count int64) error {
    // io.CopyN 会从 r 读取 count 字节，并写入 io.Discard
    // io.Discard 会丢弃所有写入的数据
    _, err := io.CopyN(io.Discard, r, count)
    if err != nil && err != io.EOF { // 忽略io.EOF错误，表示已成功读取到流末尾
        return fmt.Errorf("failed to skip %d bytes: %w", count, err)
    }
    return nil
}

func main() {
    // 示例：使用 strings.NewReader 作为 io.Reader
    data := "HEADER_1234567890_FOOTER"
    reader := strings.NewReader(data)

    fmt.Printf("原始数据流: %s\n", data)

    // 假设我们想跳过 "HEADER_" (7个字节)
    skipCount := int64(7)
    err := SkipNBytesFromReader(reader, skipCount)
    if err != nil {
        fmt.Printf("跳过字节失败: %v\n", err)
        return
    }
    fmt.Printf("成功跳过 %d 字节。\n", skipCount)

    // 读取剩余数据
    remaining, err := io.ReadAll(reader)
    if err != nil {
        fmt.Printf("读取剩余数据失败: %v\n", err)
        return
    }
    fmt.Printf("剩余数据: %s\n", string(remaining)) // 期望输出: 1234567890_FOOTER

    // 示例2: 跳过超出实际长度的字节
    reader2 := strings.NewReader("short_data")
    fmt.Printf("\n原始数据流: %s\n", "short_data")
    err = SkipNBytesFromReader(reader2, 20) // 尝试跳过20字节，但实际只有10字节
    if err != nil {
        fmt.Printf("跳过字节失败: %v\n", err) // 此时不会返回错误，因为io.CopyN在达到EOF时会返回已读取的字节数和io.EOF
    } else {
        fmt.Println("尝试跳过20字节成功 (实际可能跳过较少)。")
    }
    remaining2, _ := io.ReadAll(reader2)
    fmt.Printf("剩余数据: %s\n", string(remaining2)) // 期望输出: ""
}

type Seeker interface {
    Seek(offset int64, whence int) (int64, error)
}

package main

import (
    "fmt"
    "io"
    "strings"
)

// SkipNBytesFromReaderOptimizied 尝试从io.Reader中跳过指定数量的字节，
// 如果Reader是io.Seeker，则使用Seek方法，否则回退到io.CopyN。
func SkipNBytesFromReaderOptimizied(r io.Reader, count int64) error {
    if count < 0 {
        return fmt.Errorf("skip count cannot be negative: %d", count)
    }

    // 尝试进行类型断言，看Reader是否也实现了io.Seeker接口
    if seeker, ok := r.(io.Seeker); ok {
        // 如果是io.Seeker，则使用Seek方法进行定位
        _, err := seeker.Seek(count, io.SeekCurrent)
        if err != nil {
            return fmt.Errorf("failed to seek %d bytes: %w", count, err)
        }
        return nil
    }

    // 如果不是io.Seeker，则回退到使用io.CopyN和io.Discard
    _, err := io.CopyN(io.Discard, r, count)
    if err != nil && err != io.EOF {
        return fmt.Errorf("failed to skip %d bytes using CopyN: %w", count, err)
    }
    return nil
}

func main() {
    // 示例1：使用 strings.NewReader (实现了io.Seeker)
    data := "HEADER_1234567890_FOOTER"
    reader1 := strings.NewReader(data)

    fmt.Printf("原始数据流 (strings.Reader): %s\n", data)

    skipCount1 := int64(7) // 跳过 "HEADER_"
    err := SkipNBytesFromReaderOptimizied(reader1, skipCount1)
    if err != nil {
        fmt.Printf("跳过字节失败: %v\n", err)
        return
    }
    fmt.Printf("成功跳过 %d 字节。\n", skipCount1)

    remaining1, _ := io.ReadAll(reader1)
    fmt.Printf("剩余数据: %s\n", string(remaining1)) // 期望输出: 1234567890_FOOTER

    // 示例2：使用 io.LimitReader (未实现io.Seeker)
    // LimitReader 包装了 strings.NewReader，但它本身不是 Seeker
    limitedReader := io.LimitReader(strings.NewReader(data), int64(len(data)))

    fmt.Printf("\n原始数据流 (io.LimitReader): %s\n", data)

    skipCount2 := int64(7)
    err = SkipNBytesFromReaderOptimizied(limitedReader, skipCount2)
    if err != nil {
        fmt.Printf("跳过字节失败: %v\n", err)
        return
    }
    fmt.Printf("成功跳过 %d 字节。\n", skipCount2)

    remaining2, _ := io.ReadAll(limitedReader)
    fmt.Printf("剩余数据: %s\n", string(remaining2)) // 期望输出: 1234567890_FOOTER
}