make parase easier and allow tcp test allocate memory more than 128k

Signed-off-by: austinhmh <2287728420@qq.com>
This commit is contained in:
austinhmh 2024-12-12 16:45:07 +08:00 committed by leonrayang
parent 65726f617b
commit e095fbc2f7
4 changed files with 105 additions and 43 deletions

View File

@ -1,6 +1,7 @@
package main
import (
"flag"
"net"
"rdma_test/common"
"rdma_test/rdma"
@ -12,7 +13,7 @@ var Config = &rdma.RdmaEnvConfig{}
func ReadBytes(conn net.Conn, buf []byte) error {
offset := 0
for offset < len(buf) {
for offset < len(buf) {
n, err := conn.Read(buf[offset:])
if n == -1 || (err != nil) {
println("ReadBytes failed, err: ", err)
@ -23,8 +24,19 @@ func ReadBytes(conn net.Conn, buf []byte) error {
return nil
}
func init() {
flag.StringVar(&Config.RdmaPort, "rdma-port", "9000", "rdma-port")
flag.IntVar(&Config.MemBlockNum, "memory-block-num", 0, "memory-block-num")
flag.IntVar(&Config.MemBlockSize, "memory-block-size", 0, "memory-block-size")
flag.IntVar(&Config.MemPoolLevel, "memory-pool-level", 0, "memory-pool-level")
flag.IntVar(&Config.ConnDataSize, "connect-data-size", 0, "connect-data-size")
flag.IntVar(&Config.WqDepth, "wq-depth", 0, "wq-depth")
flag.BoolVar(&Config.EnableRdmaLog, "enable-rdma-log", false, "enable-rdma-log")
flag.IntVar(&Config.WorkerNum, "worker-num", 0, "worker-num")
flag.StringVar(&Config.RdmaLogDir, "rdma-log-dir", "", "rdma-log-dir")
}
func testRdma() {
func testRdma() {
if err := rdma.InitPool(Config); err != nil {
println("init rdma pool failed")
return
@ -34,14 +46,14 @@ func testRdma() {
go func() {
conn := &rdma.Connection{}
if err := conn.Dial(common.GParam.Ip,common.GParam.Port); err != nil {
if err := conn.Dial(common.GParam.Ip, common.GParam.Port); err != nil {
println("client rdma conn dial failed")
return
}
defer conn.Close()
for !exit {
beginTm := time.Now()
p := common.NewWritePacket(common.NormalExtentType, conn)
p := common.NewWritePacket(common.NormalExtentType, conn, uint32(common.GParam.IoSize))
p.Size = uint32(common.GParam.IoSize)
if err := p.WriteToRDMAConn(conn); err != nil {
@ -61,13 +73,13 @@ func testRdma() {
println(err.Error())
break
}
err = rdma.ReleaseDataBuffer(conn, p.RdmaBuffer, uint32(105 + common.GParam.IoSize))
err = rdma.ReleaseDataBuffer(conn, p.RdmaBuffer, uint32(105+common.GParam.IoSize))
if err != nil {
println(err.Error())
break
}
common.Stat().AddSumTime(common.GParam.IoSize, time.Now().UnixNano() / 1000 - beginTm.UnixNano() / 1000)
common.Stat().AddSumTime(common.GParam.IoSize, time.Now().UnixNano()/1000-beginTm.UnixNano()/1000)
}
}()
}
@ -78,7 +90,7 @@ func testTcp() {
for i := 0; i < common.GParam.IoDeep; i++ {
go func() {
c, err := net.Dial("tcp", common.GParam.Ip + ":" + common.GParam.Port)
c, err := net.Dial("tcp", common.GParam.Ip+":"+common.GParam.Port)
if err != nil {
println(err.Error())
return
@ -89,7 +101,7 @@ func testTcp() {
defer conn.Close()
for !exit {
beginTm := time.Now()
p := common.NewWritePacket(common.NormalExtentType, conn)
p := common.NewWritePacket(common.NormalExtentType, conn, uint32(common.GParam.IoSize))
p.Size = uint32(common.GParam.IoSize)
if err := p.WriteToConn(conn); err != nil {
@ -103,14 +115,13 @@ func testTcp() {
break
}
common.Stat().AddSumTime(common.GParam.IoSize, time.Now().UnixNano() / 1000 - beginTm.UnixNano() / 1000)
common.Stat().AddSumTime(common.GParam.IoSize, time.Now().UnixNano()/1000-beginTm.UnixNano()/1000)
}
}()
}
}
func main() {
func main() {
common.ParseParam()
if common.GParam.Protocol == "rdma" {
go testRdma()
@ -123,4 +134,3 @@ func main() {
common.Stat().Print()
}
}

View File

@ -2,10 +2,12 @@ package common
import (
"fmt"
"math"
"rdma_test/common/context"
"rdma_test/common/rate"
"sync"
"sync/atomic"
"time"
)
const (
@ -74,6 +76,38 @@ func NewNormalBufferPool() *sync.Pool {
}
}
func NewBinaryPool() *BinaryPool {
pools := make([]*sync.Pool, 32)
for i := 0; i < 32; i++ {
pools[i] = &sync.Pool{
New: func() interface{} {
buff := make([]byte, 1<<i)
return &buff
},
}
}
return &BinaryPool{pools: pools}
}
func (bp *BinaryPool) Get(size int) []byte {
var buff *[]byte
l := int(math.Ceil(math.Log2(float64(size))))
for buff == nil {
buff = bp.pools[l].Get().(*[]byte)
time.Sleep(10 * time.Millisecond)
}
return (*buff)[:size]
}
func (bp *BinaryPool) Put(buff []byte) {
l := int(math.Ceil(math.Log2(float64(cap(buff)))))
bp.pools[l].Put(&buff)
}
type BinaryPool struct {
pools []*sync.Pool
}
// BufferPool defines the struct of a buffered pool with 4 objects.
type BufferPool struct {
headPools []chan []byte
@ -81,6 +115,7 @@ type BufferPool struct {
tinyPool *sync.Pool
headPool *sync.Pool
normalPool *sync.Pool
binaryPool *BinaryPool
}
var (
@ -101,6 +136,7 @@ func NewBufferPool() (bufferP *BufferPool) {
bufferP.tinyPool = NewTinyBufferPool()
bufferP.headPool = NewHeadBufferPool()
bufferP.normalPool = NewNormalBufferPool()
bufferP.binaryPool = NewBinaryPool()
return bufferP
}
func (bufferP *BufferPool) getHead(id uint64) (data []byte) {
@ -134,8 +170,13 @@ func (bufferP *BufferPool) Get(size int) (data []byte, err error) {
} else if size == DefaultTinySizeLimit {
atomic.AddInt64(&tinyBuffersCount, 1)
return bufferP.tinyPool.Get().([]byte), nil
} else {
buff := bufferP.binaryPool.Get(size)
if len(buff) != size {
return nil, fmt.Errorf("we only can allocate buffer which less than 1 << 31 bytes")
}
return buff, nil
}
return nil, fmt.Errorf("can only support 45 or 65536 bytes")
}
func (bufferP *BufferPool) putHead(index int, data []byte) {
@ -173,6 +214,8 @@ func (bufferP *BufferPool) Put(data []byte) {
} else if size == DefaultTinySizeLimit {
bufferP.tinyPool.Put(data)
atomic.AddInt64(&tinyBuffersCount, -1)
} else {
bufferP.binaryPool.Put(data)
}
return
}

View File

@ -20,14 +20,14 @@ const (
TB
PB
OpWrite uint8 = 0x03
OpWrite uint8 = 0x03
OpOk uint8 = 0xF0
OpOk uint8 = 0xF0
ProtoMagic uint8 = 0xFF
ProtoMagic uint8 = 0xFF
TinyExtentType = 0
NormalExtentType = 1
TinyExtentType = 0
NormalExtentType = 1
BlockCount = 1024
BlockSize = 65536 * 2
@ -41,7 +41,7 @@ const (
DefaultTinySizeLimit = 1 * MB
)
var(
var (
GRequestID = int64(1)
Buffers *BufferPool
)
@ -72,20 +72,20 @@ type Packet struct {
ReqID int64
Arg []byte // for create or append ops, the data contains the address
Data []byte
RdmaBuffer []byte
RdmaBuffer []byte
StartT int64
mesg string
HasPrepare bool
}
func NewWritePacket(storeMode int, c net.Conn) *Packet {
func NewWritePacket(storeMode int, c net.Conn, size uint32) *Packet {
p := new(Packet)
p.ReqID = GenerateRequestID()
p.Magic = ProtoMagic
p.Opcode = OpWrite
if storeMode == TinyExtentType {
if GParam.Protocol == "rdma" {
dataBuffer,_ := rdma.GetDataBuffer(DefaultTinySizeLimit + 105)
dataBuffer, _ := rdma.GetDataBuffer(DefaultTinySizeLimit + 105)
p.Arg = dataBuffer[65:105]
p.Data = dataBuffer[105:]
p.RdmaBuffer = dataBuffer
@ -95,13 +95,13 @@ func NewWritePacket(storeMode int, c net.Conn) *Packet {
}
} else {
if GParam.Protocol == "rdma" {
dataBuffer,_ := rdma.GetDataBuffer(BlockSize + 105)
dataBuffer, _ := rdma.GetDataBuffer(size + 105)
p.Arg = dataBuffer[65:105]
p.Data = dataBuffer[105:]
p.RdmaBuffer = dataBuffer
} else {
p.Arg = make([]byte, 40)
p.Data, _ = Buffers.Get(BlockSize)
p.Data, _ = Buffers.Get(int(size))
}
}
return p
@ -177,7 +177,7 @@ func (p *Packet) WriteToRDMAConn(conn *rdma.Connection) (err error) {
offset += 40
if _, err = conn.WriteExternalBuffer(p.RdmaBuffer, int(105 + p.Size)); err != nil {
if _, err = conn.WriteExternalBuffer(p.RdmaBuffer, int(105+p.Size)); err != nil {
return
}
return
@ -218,7 +218,7 @@ func (p *Packet) RecvRespFromRDMAConn(c *rdma.Connection, timeoutSec int) (err e
offset += PacketHeaderSize + 8
if p.ArgLen > 0 {
p.Arg = dataBuffer[65:65+p.ArgLen]
p.Arg = dataBuffer[65 : 65+p.ArgLen]
}
offset += 40
@ -227,7 +227,7 @@ func (p *Packet) RecvRespFromRDMAConn(c *rdma.Connection, timeoutSec int) (err e
return syscall.EBADMSG
}
size := p.Size
p.Data = dataBuffer[105:105+int(size)]
p.Data = dataBuffer[105 : 105+int(size)]
return
}
@ -290,7 +290,7 @@ func (p *Packet) ReadFromRDMAConnFromCli(conn *rdma.Connection, deadlineTime tim
}
offset += PacketHeaderSize + 8 //rdma
if p.ArgLen > 0 {
p.Arg = dataBuffer[65:65+p.ArgLen]
p.Arg = dataBuffer[65 : 65+p.ArgLen]
}
offset += 40
if p.Size < 0 {
@ -299,11 +299,10 @@ func (p *Packet) ReadFromRDMAConnFromCli(conn *rdma.Connection, deadlineTime tim
}
size := p.Size
p.Data = dataBuffer[105:105+int(size)]
p.Data = dataBuffer[105 : 105+int(size)]
return
}
func ReadFull(c net.Conn, buf *[]byte, readSize int) (err error) {
*buf = make([]byte, readSize)
_, err = io.ReadFull(c, (*buf)[:readSize])
@ -360,7 +359,6 @@ func (p *Packet) SendRespToRDMAConn(conn *rdma.Connection) (err error) {
}
}()
p.MarshalHeader(dataBuffer[0:PacketHeaderSize])
offset += PacketHeaderSize + 8
if p.ArgLen != 0 {
@ -376,4 +374,4 @@ func (p *Packet) SendRespToRDMAConn(conn *rdma.Connection) (err error) {
return
}
return
}
}

View File

@ -1,6 +1,7 @@
package main
import (
"flag"
"net"
"rdma_test/common"
"rdma_test/rdma"
@ -12,7 +13,7 @@ var Config = &rdma.RdmaEnvConfig{}
func ReadBytes(conn net.Conn, buf []byte) error {
offset := 0
for offset < len(buf) {
for offset < len(buf) {
n, err := conn.Read(buf[offset:])
if n == -1 || (err != nil) {
println("ReadBytes failed, err: ", err)
@ -23,7 +24,19 @@ func ReadBytes(conn net.Conn, buf []byte) error {
return nil
}
func testRdma() {
func init() {
flag.StringVar(&Config.RdmaPort, "rdma-port", "9000", "rdma-port")
flag.IntVar(&Config.MemBlockNum, "memory-block-num", 0, "memory-block-num")
flag.IntVar(&Config.MemBlockSize, "memory-block-size", 0, "memory-block-size")
flag.IntVar(&Config.MemPoolLevel, "memory-pool-level", 0, "memory-pool-level")
flag.IntVar(&Config.ConnDataSize, "connect-data-size", 0, "connect-data-size")
flag.IntVar(&Config.WqDepth, "wq-depth", 0, "wq-depth")
flag.BoolVar(&Config.EnableRdmaLog, "enable-rdma-log", false, "enable-rdma-log")
flag.IntVar(&Config.WorkerNum, "worker-num", 0, "worker-num")
flag.StringVar(&Config.RdmaLogDir, "rdma-log-dir", "", "rdma-log-dir")
}
func testRdma() {
if err := rdma.InitPool(Config); err != nil {
println(err.Error())
return
@ -31,8 +44,8 @@ func testRdma() {
server, _ := rdma.NewRdmaServer(common.GParam.Ip, common.GParam.Port)
defer server.Close()
for {
conn,_ := server.Accept()
for {
conn, _ := server.Accept()
go func() {
for !exit {
beginTm := time.Now()
@ -56,20 +69,19 @@ func testRdma() {
break
}
common.Stat().AddSumTime(common.GParam.IoSize, time.Now().UnixNano() / 1000 - beginTm.UnixNano() / 1000)
common.Stat().AddSumTime(common.GParam.IoSize, time.Now().UnixNano()/1000-beginTm.UnixNano()/1000)
}
conn.Close()
}()
}
}
func testTcp() {
common.InitBufferPool(0)
server, _ := net.Listen("tcp", common.GParam.Ip + ":" + common.GParam.Port)
server, _ := net.Listen("tcp", common.GParam.Ip+":"+common.GParam.Port)
defer server.Close()
for {
for {
c, _ := server.Accept()
conn, _ := c.(*net.TCPConn)
conn.SetKeepAlive(true)
@ -93,14 +105,13 @@ func testTcp() {
break
}
common.Stat().AddSumTime(common.GParam.IoSize, time.Now().UnixNano() / 1000 - beginTm.UnixNano() / 1000)
common.Stat().AddSumTime(common.GParam.IoSize, time.Now().UnixNano()/1000-beginTm.UnixNano()/1000)
}
}()
}
}
func main() {
func main() {
common.ParseParam()
if common.GParam.Protocol == "rdma" {
go testRdma()