splice

gistfile1.go

package main

import (
	"flag"
	"net"
	"os"
	"syscall"

	"github.com/golang/glog"
)

func main() {
	flag.Parse()
	ln, err := net.Listen("tcp", ":8080")
	if err != nil {
		glog.Fatalf("could not listen")
	}
	glog.Info("listening")
	for {
		conn, err := ln.Accept()
		if err != nil {
			glog.Error("error accepting connection")
			continue
		}
		go handleConnection(conn, zeroCopy)
	}
}

type spliceable interface {
	File() (*os.File, error)
}

func handleConnection(iconn net.Conn, copyFunc func(in *os.File, ou *os.File)) {
	in, ok := iconn.(spliceable)
	if !ok {
		glog.Errorf("not spliceable: %+v", iconn)
		iconn.Close()
		return
	}

	ifile, err := in.File()
	iconn.Close()
	if err != nil {
		glog.Error("yolo out: %v", err)
		return
	}

	oconn, err := net.Dial("tcp", "localhost:80")
	if err != nil {
		glog.Errorf("error getting file from spliceable")
		return
	}

	out, ok := oconn.(spliceable)
	if !ok {
		glog.Errorf("not spliceable: %+v", oconn)
		oconn.Close()
		return
	}
	ofile, err := out.File()
	oconn.Close()
	if err != nil {
		glog.Error("yolo out")
		return
	}

	go copyFunc(ifile, ofile)
	go copyFunc(ofile, ifile)
}

const txLength = 4098

const (
	SPLICE_F_NONBLOCK = 0x02
)

func zeroCopy(in *os.File, out *os.File) {
	filedes := make([]int, 2)
	err := syscall.Pipe(filedes)
	if err != nil {
		glog.Errorf("zero copy create pipe error: %v", err)
		return
	}
	defer syscall.Close(filedes[0])
	defer syscall.Close(filedes[1])

	inFd := int(in.Fd())
	outFd := int(out.Fd())
	defer func() {
		syscall.SetNonblock(int(inFd), true)
		syscall.Close(inFd)
	}()
	defer func() {
		syscall.SetNonblock(int(outFd), true)
		syscall.Close(inFd)
	}()
	txLength := 2048
	for {
		_, err := syscall.Splice(inFd, nil, filedes[1], nil, txLength, SPLICE_F_NONBLOCK)
		if err != nil {
			// if err == syscall.EAGAIN {
			// 	runtime.Gosched()
			// 	continue
			// }
			glog.Errorf("error splicing into pipe: %v", err)
			return
		}
		sizeOut, err := syscall.Splice(filedes[0], nil, outFd, nil, txLength, SPLICE_F_NONBLOCK)
		if err != nil {
			// if err == syscall.EAGAIN {
			// 	glog.Infof("caught")
			// 	runtime.Gosched()
			// 	return
			// }
			glog.Errorf("error splicing out of pipe: %+v", err)
			return
		}
		if sizeOut == 0 {
			glog.Info("done")
			return
		}
	}
}

func copy(in *os.File, out *os.File) {
	defer in.Close()
	defer out.Close()
	inFd := int(in.Fd())
	outFd := int(out.Fd())
	txLength := 2048
	buffer := make([]byte, txLength)
	for {
		_, err := syscall.Read(inFd, buffer)
		if err != nil {
			glog.Errorf("error reading: %v", err)
			return
		}
		bytesWritten, err := syscall.Write(outFd, buffer)
		if err != nil {
			glog.Errorf("error writing: %v", err)
			return
		}
		if bytesWritten < txLength {
			glog.Info("done")
			return
		}
	}
}