io.Copyを調べてみた

概要

こちらを読んでそういえばio.Copyちゃんと見たことなかったなと思い調べてみました。こちらのサイトio#Copyを確認するとwriter(dst)とreader(src)を引数で受け取ってEOFに到達するかエラーが発生するまでwriteにコピーする関数。

Copy copies from src to dst until either EOF is reached on src or an error occurs. It returns the number of bytes copied and the first error encountered while copying, if any.

最後に注意書きがあり、srcがWriteToインタフェースを実装してる場合はこちらを使用するそうでなければdstにReadFromインタフェースを実装している場合はそちらが呼ばれると書かれている。（特殊対応されているのでこちらが実装されているもののほうがおそらく早くて軽いはず）

If src implements the WriterTo interface, the copy is implemented by calling src.WriteTo(dst). Otherwise, if dst implements the ReaderFrom interface, the copy is implemented by calling dst.ReadFrom(src).

定義

io.Copy関数の定義

func Copy(dst Writer, src Reader) (written int64, err error)

Writer Interfaceの定義

type Writer interface {
	Write(p []byte) (n int, err error)
}

Reader Interfaceの定義

type Reader interface {
	Read(p []byte) (n int, err error)
}

io.Copyの例

本家のサイトに例が乗っている。この場合、Readerインタフェースを実装しているstrings.NewReaderから返るReaderをsrcへ渡して、Writerインタフェースを実装しているos.Stdoutをdstに渡している。

package main

import (
	"io"
	"log"
	"os"
	"strings"
)

func main() {
	r := strings.NewReader("some io.Reader stream to be read\n")

	if _, err := io.Copy(os.Stdout, r); err != nil {
		log.Fatal(err)
	}

}

io.Copyの基本的なフロー

WriterToやReaderFrom、LimitedReaderなどが実装されていないシンプルなフローを追う

io.Copy

内部でcopyBuffer関数を呼んでいる

func Copy(dst Writer, src Reader) (written int64, err error) {
	return copyBuffer(dst, src, nil)
}

copyBuffer

copyBuffer(dst, src, nil)の最後の引数はbufでですがnilで呼び出されています。bufがnilの場合は32kbのサイズのbufを生成します。（このサイズが小さいか大きいかは用途によると思うので指定したい場合はCopyBufferを使用して明示的に指定する方法をとった方がいい）

if buf == nil {
	size := 32 * 1024
	buf = make([]byte, size)
}

srcからbufferに読み込み。読み込みがあれば、bufferをdstへ書き込みない場合は、EOFかをチェックして終了。

for {
	nr, er := src.Read(buf)
	if nr > 0 {
		nw, ew := dst.Write(buf[0:nr])
		if nw < 0 || nr < nw {
			nw = 0
			if ew == nil {
				ew = errInvalidWrite
			}
		}
		written += int64(nw)
		if ew != nil {
			err = ew
			break
		}
		if nr != nw {
			err = ErrShortWrite
			break
		}
	}
	if er != nil {
		if er != EOF {
			err = er
		}
		break
	}
}

srcがWriterToを実装していた場合

io#WriteStringのコメントを見ると、allocationとcopyを避けられると書かれている。（なので、通常フローよりも負荷が低いはず

// If the reader has a WriteTo method, use it to do the copy.
// Avoids an allocation and a copy.

WriteToが実装されていれば、dstを引数にとってそのままWriteToを実行して終了する

if wt, ok := src.(WriterTo); ok {
	return wt.WriteTo(dst)
}

WriterTo interfaceの定義

type WriterTo interface {
	WriteTo(w Writer) (n int64, err error)
}

実際にこちらを実装しているライブラリ

 ag 'WriteTo implements' -l
bufio/bufio.go
net/net.go
net/iprawsock.go
net/unixsock.go
net/udpsock.go
strings/reader.go
bytes/reader.go

strings.Reader.WriteTo

strings.ReaderはWriteToを実装しているので内部ではこちらが呼ばれます。
最終的にio.WriteStringを呼んで書き込んでいます

func (r *Reader) WriteTo(w io.Writer) (n int64, err error) {
	r.prevRune = -1
	if r.i >= int64(len(r.s)) {
		return 0, nil
	}
	s := r.s[r.i:]
	m, err := io.WriteString(w, s)
	if m > len(s) {
		panic("strings.Reader.WriteTo: invalid WriteString count")
	}
	r.i += int64(m)
	n = int64(m)
	if m != len(s) && err == nil {
		err = io.ErrShortWrite
	}
	return
}

StringWriter Interface

• io.WriteStringはWriter側にStringWriterが実装されていた場合
• そちらで処理されます
• ここで渡されているWriterはos.Fileです
  • https://pkg.go.dev/os@go1.18.1#File
• os.FileはStringWriterインタフェースを満たしているので
  • https://pkg.go.dev/io#StringWriter
• WriterStringメソッドを使用して処理を行います

func WriteString(w Writer, s string) (n int, err error) {
	if sw, ok := w.(StringWriter); ok {
		return sw.WriteString(s)
	}
	return w.Write([]byte(s))
}

File.WriteString

• WriteStringの実装は以下
• bytesスライスを宣言してunsafe.Pointerを得て
• unsafeheader.Sliceへキャストしています
  • https://pkg.go.dev/golang.org/x/sys/internal/unsafeheader#Slice
• 受け取った文字列からデータを受取直接
• データ構造へ値を渡してbyteへ変換を行いWriteしています

func (f *File) WriteString(s string) (n int, err error) {
	var b []byte
	hdr := (*unsafeheader.Slice)(unsafe.Pointer(&b))
	hdr.Data = (*unsafeheader.String)(unsafe.Pointer(&s)).Data
	hdr.Cap = len(s)
	hdr.Len = len(s)
	return f.Write(b)
}

dstがReaderFromを実装してる場合

• もとの処理ではこちらのフローに入らないので
• os.Fileを渡すか別途Readerを用意する必要がある
• 今回は独自のReader Interfaceを満たす構造体を用意した

package main

import (
	"io"
	"log"
	"os"
)

type myReader struct {
}

func (r *myReader) Read(p []byte) (n int, err error) {
	copy(p, []byte("hello,world\n"))
	return len(p), io.EOF
}

func main() {
	r := &myReader{}
	if _, err := io.Copy(os.Stdout, r); err != nil {
		log.Fatal(err)
	}
}

• 途中で*Fileで判定している箇所があり
• 他のメソッドを足せば行けるとは思うんですが
• 面倒になったのでファイルを用意して開きました

package main

import (
	"io"
	"log"
	"os"
)

func main() {
	r, err := os.Open("hoge.txt")
	if err != nil {
		panic(err)
	}

	if _, err := io.Copy(os.Stdout, r); err != nil {
		log.Fatal(err)
	}
}

hello,world

ReaderFrom interface

type ReaderFrom interface {
	ReadFrom(r Reader) (n int64, err error)
}

• 実装してるモジュールは以下

 ag 'ReadFrom implements' -l    
bufio/bufio.go
net/tcpsock.go
net/iprawsock.go
net/unixsock.go
net/udpsock.go
os/file.go

File.ReadFrom

• File.ReadFromの中でreadFromを呼んでいます
• このあたりから_linuxというファイルに移動しています

// ReadFrom implements io.ReaderFrom.
func (f *File) ReadFrom(r io.Reader) (n int64, err error) {
	n, handled, e := f.readFrom(r)
	return n, f.wrapErr("write", e)
}

• readFromは更にpollCopyFileRangeという関数を呼んでいます

func (f *File) readFrom(r io.Reader) (written int64, handled bool, err error) {
	written, handled, err = pollCopyFileRange(&f.pfd, &src.pfd, remain)
	if lr != nil {
		lr.N -= written
	}
	return written, handled, NewSyscallError("copy_file_range", err)
}

• pollCopyFileRangeはpoll.CopyFileRangeが実態

var pollCopyFileRange = poll.CopyFileRange

• 追記
• 某Slackできいてきたのですが
• このようにpackage変数に代入する理由は
• テストでの置き換えの為だそうです
  • https://github.com/golang/go/blob/d5de62df152baf4de6e9fe81933319b86fd95ae4/src/os/readfrom_linux_test.go#L407

CopyFileRange

• atomic.LoadInt32をサポートしているかをチェック
• サポートしてない場合はカーネルバージョンをチェックして
• atomic.StoreInt32が使えるならこちらを使用
• 現状のシステムがKernel5.15なのでこちらのフローに入ります
• https://github.com/golang/go/blob/d5de62df152baf4de6e9fe81933319b86fd95ae4/src/os/readfrom_linux_test.go#L407

func CopyFileRange(dst, src *FD, remain int64) (written int64, handled bool, err error) {
	if supported := atomic.LoadInt32(&copyFileRangeSupported); supported == 0 {
		return 0, false, nil
	} else if supported == -1 {
		major, minor := kernelVersion()
		if major > 5 || (major == 5 && minor >= 3) {
			atomic.StoreInt32(&copyFileRangeSupported, 1)
		}
	}
	for remain > 0 {
		max := remain
		n, err := copyFileRange(dst, src, int(max))
	}
	return written, true, nil
}

copyFileRange

• Lockなどをかけたあとに
• unix.CopyFileRangeを呼び出しています

func copyFileRange(dst, src *FD, max int) (written int64, err error) {
	if err := dst.writeLock(); err != nil {
		return 0, err
	}
	defer dst.writeUnlock()
	if err := src.readLock(); err != nil {
		return 0, err
	}
	defer src.readUnlock()
	var n int
	for {
		n, err = unix.CopyFileRange(src.Sysfd, nil, dst.Sysfd, nil, max, 0)
		if err != syscall.EINTR {
			break
		}
	}
	return int64(n), err
}

unix.CopyFileRange

• ここで直接Syscall6を呼び出して書き込んでいます
• copyFileRangeTrapは326と定義されています

func CopyFileRange(rfd int, roff *int64, wfd int, woff *int64, len int, flags int) (n int, err error) {
	r1, _, errno := syscall.Syscall6(copyFileRangeTrap,
		uintptr(rfd),
		uintptr(unsafe.Pointer(roff)),
		uintptr(wfd),
		uintptr(unsafe.Pointer(woff)),
		uintptr(len),
		uintptr(flags),
	)
	n = int(r1)
	if errno != 0 {
		err = errno
	}
	return
}

syscall.Syscall6

• アセンブラが書かれています
• runtime.entersyscallが呼ばれているようです
• ココらへんの詳細に関しては後日見ます
• こちらのリンクが参考になります
  • https://zenn.dev/hsaki/books/golang-concurrency/viewer/internalcase
  • https://ascii.jp/elem/000/001/267/1267477/

TEXT ·Syscall6(SB),NOSPLIT,$0-80
	CALL	runtime·entersyscall(SB)
	MOVQ	a1+8(FP), DI
	MOVQ	a2+16(FP), SI
	MOVQ	a3+24(FP), DX
	MOVQ	a4+32(FP), R10
	MOVQ	a5+40(FP), R8
	MOVQ	a6+48(FP), R9
	MOVQ	trap+0(FP), AX	// syscall entry
	SYSCALL
	CMPQ	AX, $0xfffffffffffff001
	JLS	ok6
	MOVQ	$-1, r1+56(FP)
	MOVQ	$0, r2+64(FP)
	NEGQ	AX
	MOVQ	AX, err+72(FP)
	CALL	runtime·exitsyscall(SB)
	RET
ok6:
	MOVQ	AX, r1+56(FP)
	MOVQ	DX, r2+64(FP)
	MOVQ	$0, err+72(FP)
	CALL	runtime·exitsyscall(SB)
	RET

参考

• https://itkq.jp/blog/2017/05/10/linux-file-and-io/
• https://zenn.dev/nobonobo/articles/297dc5cbc554d6