用于测试golang的网络单链接性能

calcSpeed_linux.go

// +build linux
package main

type TransferSpeed struct{}

var RSpeed *TransferSpeed = NewTransferSpeed()
var WSpeed *TransferSpeed = NewTransferSpeed()

func UpdateReadIO(bytes uint32) {
    return
}

func UpdateWriteIO(bytes uint32) {
	return
}

func NewTransferSpeed() *TransferSpeed {
	return &TransferSpeed{}
}

calcSpeed_windows.go

// +build windows
package main

import (
    "fmt"
	"os"
	"sync"
	"syscall"
	"time"
	"unsafe"
)

var (
	modKernel32 = syscall.NewLazyDLL("kernel32.dll")

	procQueryPerformanceFrequency = modKernel32.NewProc("QueryPerformanceFrequency")
	procQueryPerformanceCounter   = modKernel32.NewProc("QueryPerformanceCounter")
)

func QueryPerformaceFrequency() (frequency int64, err error) {
	_, _, e0 := syscall.Syscall(procQueryPerformanceFrequency.Addr(), 1, uintptr(unsafe.Pointer(&frequency)), 0, 0)
	if e0 != 0 {
		err = error(e0)
	}
	return
}

func QueryPerformanceCounter() (tick int64, err error) {
	_, _, e0 := syscall.Syscall(procQueryPerformanceCounter.Addr(), 1, uintptr(unsafe.Pointer(&tick)), 0, 0)
	if e0 != 0 {
		err = error(e0)
	}
	return
}

type TransferSpeed struct {
	sumTransferBytes uint32
	speed            float64
	lastTick         uint32
	maxElaspeMs      uint32
	sampleUnit       uint32
	frequency        uint32
	lock             sync.Mutex
}

var RSpeed *TransferSpeed = NewTransferSpeed()
var WSpeed *TransferSpeed = NewTransferSpeed()

func UpdateReadIO(bytes uint32) {
	RSpeed.UpdateIOBytes(bytes)
}

func UpdateWriteIO(bytes uint32) {
	WSpeed.UpdateIOBytes(bytes)
}

func NewTransferSpeed() *TransferSpeed {
	frequency, _ := QueryPerformaceFrequency()
	return &TransferSpeed{
		frequency:   uint32(frequency / 1000),
		sampleUnit:  uint32(frequency / 100000),
		maxElaspeMs: 1024,
	}
}

func (t *TransferSpeed) getCurrentTick() uint32 {
	tick, _ := QueryPerformanceCounter()
	return uint32(tick) / t.frequency
}

func (t *TransferSpeed) getElaspe(currentTick uint32) uint32 {
	elaspeMs := currentTick - t.lastTick
	if elaspeMs == 0 {
		elaspeMs = 10
	}
	return elaspeMs
}

func (t *TransferSpeed) GetSpeedRate(currentTick uint32) float64 {
	t.lock.Lock()
	defer t.lock.Unlock()
	if currentTick == 0 {
		currentTick = t.getCurrentTick()
	}
	elaspeMs := t.getElaspe(currentTick)
	speed := float64(t.sumTransferBytes) / float64(elaspeMs) * 1000
	if elaspeMs >= t.maxElaspeMs {
		return speed
	}
	rate := t.speed*float64(t.maxElaspeMs-elaspeMs)/float64(t.maxElaspeMs) + speed*float64(elaspeMs)/float64(t.maxElaspeMs)
	return rate
}

func (t *TransferSpeed) UpdateIOBytes(bytes uint32) {
	t.lock.Lock()
	currentTick := t.getCurrentTick()
	elaspeMs := t.getElaspe(currentTick)

	t.sumTransferBytes += bytes
	if elaspeMs >= t.sampleUnit {
		t.lock.Unlock()
		t.speed = t.GetSpeedRate(currentTick)
		t.lock.Lock()
		t.sumTransferBytes = 0
		t.lastTick = currentTick
	}
	t.lock.Unlock()
}

func init() {
	go func() {
		ticker := time.NewTicker(time.Duration(1) * time.Second)
		for {
			select {
			case <-ticker.C:
				{
					fmt.Fprintln(os.Stdout, "Read  Speed ", uint32(RSpeed.GetSpeedRate(0)/1024/1024), " MB/S")
					fmt.Fprintln(os.Stdout, "Write Speed ", uint32(WSpeed.GetSpeedRate(0)/1024/1024), " MB/S")
				}
			}
		}
	}()
}

fakehttp.go

package main

import (
    "io"
	"io/ioutil"
	"net/http"
	"fmt"
)

func fakeHTTPServer() error {
	http.HandleFunc("/test", func(response http.ResponseWriter, request *http.Request){
		response.Header().Add("Content-Length", fmt.Sprintf("%d", 1023 * 1024 * 1024))
		buff := make([]byte, 1024 * 64)
	    totalLength := 1024 * 1024 * 1024
		for totalLength > 0 {
			_, err := response.Write(buff)
			if err != nil {
				handleError(err)
				break
			}
			totalLength -= 64 * 1024
		}
	})

	return http.ListenAndServe(":8888", nil)
}

func fakseHTTPClient(url string) error {
	response, err := http.Get(url)
	if err != nil {
		return err
	}
	defer response.Body.Close()
	_, err = io.Copy(ioutil.Discard, response.Body)
	return err
}

main.go

package main

import (
    "flag"
	"fmt"
	"io"
	"net"
	"os"
	"runtime"
)

const buffLength = 64 * 1024

func handleError(err error) {
	if err != nil {
		fmt.Fprintln(os.Stderr, "err:", err.Error())
	}
}

func readBuff(conn io.Reader, buff []byte) error {
	totalLen := len(buff)
	pos := 0
	for pos < totalLen {
		n, err := conn.Read(buff[pos:totalLen])
		if err != nil {
			if err != io.EOF {
				return err
			}
			break
		}
		pos += n
	}
	return nil
}

func writeBuff(conn io.Writer, buff []byte) error {
	totalLen := len(buff)
	pos := 0
	for pos < totalLen {
		n, err := conn.Write(buff[pos:totalLen])
		if err != nil {
			return err
		}
		pos += n
	}
	return nil
}

func asyncRecv(conn io.Reader, bufSize int, closure func()) {
	defer closure()
	buff := make([]byte, bufSize, bufSize)
	for {
		err := readBuff(conn, buff)
		if err != nil {
			handleError(err)
			break
		}
		UpdateReadIO(uint32(len(buff)))
		closure()
	}
}

func runClient(svrIp string) {
	conn, err := net.DialTCP("tcp4", &net.TCPAddr{}, &net.TCPAddr{IP: net.ParseIP(svrIp), Port: 9090})
	if err != nil {
		handleError(err)
		return
	}
	defer conn.Close()
	conn.SetNoDelay(true)

	// When recv success, notify goroutine to request remained data.
	writeSig := make(chan bool)
	// async recv. TCP/IP is send/recv double channels.
	go asyncRecv(conn, buffLength, func() {
		writeSig <- true
	})
	// fake as a client downloader. First send data to request download.
	buff := make([]byte, 1024, 1024)
	for {
		err := writeBuff(conn, buff)
		if err != nil {
			handleError(err)
			break
		}
		UpdateWriteIO(uint32(len(buff)))
		// Wait for recv file's data success. So iterator download remaining data.
		<-writeSig
	}
}

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())
	var svrIp string
	var svrUrl string
	var isFakeHttpSvr bool
	flag.StringVar(&svrUrl, "url", "", "-url=http://127.0.0.1:8888/test")
	flag.BoolVar(&isFakeHttpSvr, "fakeHttp", false, "-fakeHttp=true")
	flag.StringVar(&svrIp, "ip", "", "-ip=1.1.1.1")
	flag.Parse()
	if len(svrIp) != 0 {
		runClient(svrIp)
		return
	} else if len(svrUrl) != 0 {
		fakseHTTPClient(svrUrl)
		return
	} else if isFakeHttpSvr {
		fakeHTTPServer()
		return
	}

	listener, err := net.ListenTCP("tcp4", &net.TCPAddr{Port: 9090})
	if err != nil {
		handleError(err)
		return
	}
	defer listener.Close()
	for {
		conn, err := listener.AcceptTCP()
		if err != nil {
			handleError(err)
			break
		}
		fmt.Fprintln(os.Stdout, "new client come in!, ip=", conn.RemoteAddr().String())
		conn.SetNoDelay(true)
		go func(conn net.Conn) {
			defer fmt.Fprintln(os.Stdout, "one client disconnect!")
			defer conn.Close()
			// When recv success, notify to request other data.
			writeSig := make(chan bool)
			// Async recv data. TCP/IP is double channels.
			go asyncRecv(conn, 1024, func() {
				writeSig <- true
			})
			// When recv client's request, send target data.
			buff := make([]byte, buffLength, buffLength)
			for {
				// When recv success, the channel will be wakeup.
				<-writeSig
				// Write target data
				err := writeBuff(conn, buff)
				if err != nil {
					handleError(err)
					break
				}
				UpdateWriteIO(uint32(len(buff)))
			}
		}(conn)
	}
}