waf-platform/EdgeNode/internal/utils/mmap/shared_mmap_test.go

// Copyright 2024 GoEdge CDN goedge.cdn@gmail.com. All rights reserved. Official site: https://goedge.cn .
//go:build plus

package mmap_test

import (
	"bytes"
	"crypto/md5"
	"fmt"
	"github.com/TeaOSLab/EdgeNode/internal/utils/bytepool"
	"github.com/TeaOSLab/EdgeNode/internal/utils/mmap"
	"github.com/TeaOSLab/EdgeNode/internal/utils/testutils"
	_ "github.com/iwind/TeaGo/bootstrap"
	"io"
	"os"
	"sync"
	"testing"
)

func TestOpenSharedMMAP(t *testing.T) {
	for i := 0; i < 3; i++ {
		func() {
			fp, err := mmap.OpenSharedMMAP("./shared_mmap.go")
			if err != nil {
				t.Fatal(err)
			}

			defer func() {
				_ = fp.Close()
			}()

			var buf = make([]byte, fp.Size())
			n, err := fp.ReadAt(buf, 0)
			if err != nil && err != io.EOF {
				t.Fatal(err)
			}
			t.Log(len(buf[:n]), fp.Stat())
		}()
	}
}

func TestOpenSharedMMAP_Memory(t *testing.T) {
	var isSingleTesting = testutils.IsSingleTesting()

	var count = 10
	if isSingleTesting {
		count = 10_000
	}

	var stat1 = testutils.ReadMemoryStat()

	var fps []*mmap.SharedMMAP

	var buf = make([]byte, 1024)
	for i := 0; i < count; i++ {
		//if i > 0 && i%100 == 0 {
		//	t.Log(i)
		//}
		fp, openErr := mmap.OpenSharedMMAP("./shared_mmap.go")
		if openErr != nil {
			t.Fatal(openErr)
		}
		fps = append(fps, fp)

		var off int64
		for {
			n, err := fp.ReadAt(buf, off)
			if err != nil {
				if err == io.EOF {
					break
				}
				t.Fatal(err)
			}
			off += int64(n)
		}
	}

	var stat2 = testutils.ReadMemoryStat()
	t.Log((stat2.HeapInuse-stat1.HeapInuse)>>20, "MB")

	for _, fp := range fps {
		_ = fp.Close()
	}
}

func TestOpenSharedMMAP_Concurrent(t *testing.T) {
	const filename = "./shared_mmap.go"

	data, err := os.ReadFile(filename)
	if err != nil {
		t.Fatal(err)
	}

	var rawMdString string
	{
		var md = md5.New()
		md.Write(data)
		rawMdString = fmt.Sprintf("%x", md.Sum(nil))
	}
	t.Log(rawMdString)

	fp, err := mmap.OpenSharedMMAP(filename)
	if err != nil {
		t.Fatal(err)
	}

	const concurrent = 2048
	var wg = &sync.WaitGroup{}
	wg.Add(concurrent)

	for i := 0; i < concurrent; i++ {
		go func() {
			defer wg.Done()

			var buf = bytes.NewBuffer(nil)
			n, readErr := fp.Write(buf, 0)
			if readErr != nil && readErr != io.EOF {
				t.Log(readErr)
				return
			}
			var md = md5.New()
			md.Write(buf.Bytes()[:n])
			var mdString = fmt.Sprintf("%x", md.Sum(nil))
			if mdString != rawMdString {
				t.Log("md not equal")
				return
			}
		}()
	}

	wg.Wait()
}

func TestOpenSharedMMAP_Close(t *testing.T) {
	fp, err := mmap.OpenSharedMMAP("./shared_mmap.go")
	if err != nil {
		t.Fatal(err)
	}

	for i := 0; i < 10; i++ {
		_ = fp.Close()
	}
}

func TestCountSharedMMAP(t *testing.T) {
	t.Log(mmap.CountSharedMMAP())
}

func BenchmarkMMAP_Shared(b *testing.B) {
	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			func() {
				reader, err := mmap.OpenSharedMMAP("../kvstore/query.go")
				if err != nil {
					b.Fatal(err)
				}
				defer func() {
					_ = reader.Close()
				}()

				var data = make([]byte, reader.Size())
				n, readErr := reader.ReadAt(data, 0)
				if readErr != nil {
					b.Fatal(readErr)
				}
				if int64(n) != reader.Size() {
					b.Fatal("size not equal")
				}
			}()
		}
	})
}

func BenchmarkMMAP_Shared2(b *testing.B) {
	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			func() {
				reader, err := mmap.OpenSharedMMAP("../kvstore/query.go")
				if err != nil {
					b.Fatal(err)
				}
				defer func() {
					_ = reader.Close()
				}()

				var writer = &bytes.Buffer{}
				n, readErr := reader.Write(writer, 0)
				if readErr != nil {
					b.Fatal(readErr)
				}
				if int64(n) != reader.Size() {
					b.Fatal("size not equal")
				}
			}()
		}
	})
}

func BenchmarkMMAP_MMAP(b *testing.B) {
	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			func() {
				fp, err := os.Open("../kvstore/query.go")
				if err != nil {
					b.Fatal(err)
				}

				stat, err := fp.Stat()
				if err != nil {
					b.Fatal(err)
				}

				reader, err := mmap.OpenFile(fp, stat.Size())
				if err != nil {
					b.Fatal(err)
				}
				defer func() {
					_ = reader.Close()
				}()

				_, _ = reader.Write(io.Discard, 0)
			}()
		}
	})
}

func BenchmarkMMAP_FP_Read(b *testing.B) {
	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			func() {
				fp, err := os.Open("../kvstore/query.go")
				if err != nil {
					b.Fatal(err)
				}

				defer func() {
					_ = fp.Close()
				}()

				var buf = bytepool.Pool1k.Get()
				defer bytepool.Pool1k.Put(buf)

				for {
					_, readErr := fp.Read(buf.Bytes)
					if readErr != nil {
						return
					}
				}
			}()
		}
	})
}