refactor the project directory structure

refactor netlist & domain data loader

dispatcher/coremain/run.go

@@ -21,8 +21,8 @@ import (
 	"fmt"
 	"github.com/IrineSistiana/mosdns/dispatcher/handler"
 	"github.com/IrineSistiana/mosdns/dispatcher/mlog"
+	"github.com/IrineSistiana/mosdns/dispatcher/pkg/concurrent_limiter"
 	_ "github.com/IrineSistiana/mosdns/dispatcher/plugin"
-	"github.com/IrineSistiana/mosdns/dispatcher/utils"
 	"go.uber.org/zap"
 	"go.uber.org/zap/zapcore"
 	"os"
@@ -93,7 +93,7 @@ func loadConfig(f string, depth int) {
 	if n < 1 {
 		n = 1
 	}
-	pool := utils.NewConcurrentLimiter(n)
+	pool := concurrent_limiter.NewConcurrentLimiter(n)
 	wg := new(sync.WaitGroup)
 	for i, pluginConfig := range c.Plugin {
 		if len(pluginConfig.Tag) == 0 || len(pluginConfig.Type) == 0 {

dispatcher/pkg/concurrent_limiter/client_limiter.go

@@ -0,0 +1,69 @@
+//     Copyright (C) 2020-2021, IrineSistiana
+//
+//     This file is part of mosdns.
+//
+//     mosdns is free software: you can redistribute it and/or modify
+//     it under the terms of the GNU General Public License as published by
+//     the Free Software Foundation, either version 3 of the License, or
+//     (at your option) any later version.
+//
+//     mosdns is distributed in the hope that it will be useful,
+//     but WITHOUT ANY WARRANTY; without even the implied warranty of
+//     MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+//     GNU General Public License for more details.
+//
+//     You should have received a copy of the GNU General Public License
+//     along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+package concurrent_limiter
+
+import (
+	"github.com/IrineSistiana/mosdns/dispatcher/pkg/concurrent_map"
+)
+
+type ClientQueryLimiter struct {
+	maxQueries int
+	m          *concurrent_map.ConcurrentMap
+}
+
+func NewClientQueryLimiter(maxQueries int) *ClientQueryLimiter {
+	return &ClientQueryLimiter{
+		maxQueries: maxQueries,
+		m:          concurrent_map.NewConcurrentMap(64),
+	}
+}
+
+func (l *ClientQueryLimiter) Acquire(key string) bool {
+	return l.m.TestAndSet(key, l.acquireTestAndSet)
+}
+
+func (l *ClientQueryLimiter) acquireTestAndSet(v interface{}, ok bool) (newV interface{}, wantUpdate, passed bool) {
+	n := 0
+	if ok {
+		n = v.(int)
+	}
+	if n >= l.maxQueries {
+		return nil, false, false
+	}
+	n++
+	return n, true, true
+}
+
+func (l *ClientQueryLimiter) doneTestAndSet(v interface{}, ok bool) (newV interface{}, wantUpdate, passed bool) {
+	if !ok {
+		panic("ClientQueryLimiter doneTestAndSet: value is not exist")
+	}
+	n := v.(int)
+	n--
+	if n < 0 {
+		panic("ClientQueryLimiter doneTestAndSet: value becomes negative")
+	}
+	if n == 0 {
+		return nil, true, true
+	}
+	return n, true, true
+}
+
+func (l *ClientQueryLimiter) Done(key string) {
+	l.m.TestAndSet(key, l.doneTestAndSet)
+}

dispatcher/utils/server_handler_test.go → dispatcher/pkg/concurrent_limiter/client_limiter_test.go

@@ -1,4 +1,4 @@
-package utils
+package concurrent_limiter
 
 import (
 	"strconv"

dispatcher/pkg/concurrent_limiter/concurrent_limiter.go

@@ -0,0 +1,55 @@
+//     Copyright (C) 2020-2021, IrineSistiana
+//
+//     This file is part of mosdns.
+//
+//     mosdns is free software: you can redistribute it and/or modify
+//     it under the terms of the GNU General Public License as published by
+//     the Free Software Foundation, either version 3 of the License, or
+//     (at your option) any later version.
+//
+//     mosdns is distributed in the hope that it will be useful,
+//     but WITHOUT ANY WARRANTY; without even the implied warranty of
+//     MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+//     GNU General Public License for more details.
+//
+//     You should have received a copy of the GNU General Public License
+//     along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+package concurrent_limiter
+
+import "fmt"
+
+// ConcurrentLimiter is a soft limiter.
+type ConcurrentLimiter struct {
+	bucket chan struct{}
+}
+
+// NewConcurrentLimiter returns a ConcurrentLimiter, max must > 0.
+func NewConcurrentLimiter(max int) *ConcurrentLimiter {
+	if max <= 0 {
+		panic(fmt.Sprintf("ConcurrentLimiter: invalid max arg: %d", max))
+	}
+
+	bucket := make(chan struct{}, max)
+	for i := 0; i < max; i++ {
+		bucket <- struct{}{}
+	}
+
+	return &ConcurrentLimiter{bucket: bucket}
+}
+
+func (l *ConcurrentLimiter) Wait() <-chan struct{} {
+	return l.bucket
+}
+
+func (l *ConcurrentLimiter) Done() {
+	select {
+	case l.bucket <- struct{}{}:
+	default:
+		panic("ConcurrentLimiter: bucket overflow")
+	}
+}
+
+func (l *ConcurrentLimiter) Available() int {
+	return len(l.bucket)
+}

dispatcher/pkg/concurrent_limiter/concurrent_limiter_test.go

@@ -0,0 +1,34 @@
+package concurrent_limiter
+
+import (
+	"context"
+	"sync"
+	"testing"
+	"time"
+)
+
+func Test_ConcurrentLimiter_acquire_release(t *testing.T) {
+	l := NewConcurrentLimiter(500)
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second*2)
+	defer cancel()
+
+	wg := new(sync.WaitGroup)
+	wg.Add(1000)
+	for i := 0; i < 1000; i++ {
+		go func() {
+			defer wg.Done()
+			select {
+			case <-l.Wait():
+				time.Sleep(time.Millisecond * 200)
+				l.Done()
+			case <-ctx.Done():
+				t.Fail()
+			}
+		}()
+	}
+
+	wg.Wait()
+	if l.Available() != 500 {
+		t.Fatal("token leaked")
+	}
+}

dispatcher/utils/concurrent_lru.go → dispatcher/pkg/concurrent_lru/concurrent_lru.go

@@ -15,9 +15,10 @@
 //     You should have received a copy of the GNU General Public License
 //     along with this program.  If not, see <https://www.gnu.org/licenses/>.
 
-package utils
+package concurrent_lru
 
 import (
+	lru2 "github.com/IrineSistiana/mosdns/dispatcher/pkg/lru"
 	"hash/maphash"
 	"sync"
 )
@@ -40,7 +41,7 @@ func NewConcurrentLRU(
 	for i := range cl.l {
 		cl.l[i] = &shardedLRU{
 			onGet: onGet,
-			lru:   NewLRU(maxSizePerShard, onEvict),
+			lru:   lru2.NewLRU(maxSizePerShard, onEvict),
 		}
 	}
 
@@ -95,7 +96,7 @@ type shardedLRU struct {
 	onGet func(key string, v interface{}) interface{}
 
 	sync.Mutex
-	lru *LRU
+	lru *lru2.LRU
 }
 
 func (sl *shardedLRU) Add(key string, v interface{}) {

dispatcher/utils/concurrent_lru_test.go → dispatcher/pkg/concurrent_lru/concurrent_lru_test.go

@@ -1,4 +1,4 @@
-package utils
+package concurrent_lru
 
 import (
 	"reflect"

dispatcher/utils/concurrent_map.go → dispatcher/pkg/concurrent_map/concurrent_map.go

@@ -15,7 +15,7 @@
 //     You should have received a copy of the GNU General Public License
 //     along with this program.  If not, see <https://www.gnu.org/licenses/>.
 
-package utils
+package concurrent_map
 
 import (
 	"hash/maphash"

dispatcher/utils/concurrent_map_test.go → dispatcher/pkg/concurrent_map/concurrent_map_test.go

@@ -1,4 +1,4 @@
-package utils
+package concurrent_map
 
 import (
 	"strconv"

dispatcher/utils/net_io.go → dispatcher/pkg/dnsutils/net_io.go

@@ -15,11 +15,12 @@
 //     You should have received a copy of the GNU General Public License
 //     along with this program.  If not, see <https://www.gnu.org/licenses/>.
 
-package utils
+package dnsutils
 
 import (
 	"encoding/binary"
 	"fmt"
+	"github.com/IrineSistiana/mosdns/dispatcher/pkg/pool"
 	"github.com/miekg/dns"
 	"io"
 	"net"
@@ -60,8 +61,8 @@ func (e *IOErr) Unwrap() error {
 // An io err will be wrapped into an IOErr.
 // IsIOErr(err) can check and unwrap the inner io err.
 func ReadUDPMsgFrom(c net.PacketConn, bufSize int) (m *dns.Msg, from net.Addr, n int, err error) {
-	buf := GetMsgBuf(bufSize)
-	defer ReleaseMsgBuf(buf)
+	buf := pool.GetMsgBuf(bufSize)
+	defer pool.ReleaseMsgBuf(buf)
 
 	n, from, err = c.ReadFrom(buf)
 	if err != nil {
@@ -84,8 +85,8 @@ func ReadUDPMsgFrom(c net.PacketConn, bufSize int) (m *dns.Msg, from net.Addr, n
 
 // ReadMsgFromUDP See ReadUDPMsgFrom.
 func ReadMsgFromUDP(c io.Reader, bufSize int) (m *dns.Msg, n int, err error) {
-	buf := GetMsgBuf(bufSize)
-	defer ReleaseMsgBuf(buf)
+	buf := pool.GetMsgBuf(bufSize)
+	defer pool.ReleaseMsgBuf(buf)
 
 	n, err = c.Read(buf)
 	if err != nil {
@@ -108,11 +109,11 @@ func ReadMsgFromUDP(c io.Reader, bufSize int) (m *dns.Msg, n int, err error) {
 // An io err will be wrapped into an IOErr.
 // IsIOErr(err) can check and unwrap the inner io err.
 func WriteMsgToUDP(c io.Writer, m *dns.Msg) (n int, err error) {
-	mRaw, buf, err := PackBuffer(m)
+	mRaw, buf, err := pool.PackBuffer(m)
 	if err != nil {
 		return 0, err
 	}
-	defer ReleaseMsgBuf(buf)
+	defer pool.ReleaseMsgBuf(buf)
 
 	return WriteRawMsgToUDP(c, mRaw)
 }
@@ -128,11 +129,11 @@ func WriteRawMsgToUDP(c io.Writer, b []byte) (n int, err error) {
 
 // WriteUDPMsgTo See WriteMsgToUDP.
 func WriteUDPMsgTo(m *dns.Msg, c net.PacketConn, to net.Addr) (n int, err error) {
-	mRaw, buf, err := PackBuffer(m)
+	mRaw, buf, err := pool.PackBuffer(m)
 	if err != nil {
 		return 0, err
 	}
-	defer ReleaseMsgBuf(buf)
+	defer pool.ReleaseMsgBuf(buf)
 
 	n, err = c.WriteTo(mRaw, to)
 	if err != nil {
@@ -161,8 +162,8 @@ func ReadMsgFromTCP(c io.Reader) (m *dns.Msg, n int, err error) {
 		return nil, n, dns.ErrShortRead
 	}
 
-	buf := GetMsgBuf(int(length))
-	defer ReleaseMsgBuf(buf)
+	buf := pool.GetMsgBuf(int(length))
+	defer pool.ReleaseMsgBuf(buf)
 
 	n2, err := io.ReadFull(c, buf)
 	n = n + n2
@@ -185,11 +186,11 @@ func ReadMsgFromTCP(c io.Reader) (m *dns.Msg, n int, err error) {
 // An io err will be wrapped into an IOErr.
 // IsIOErr(err) can check and unwrap the inner io err.
 func WriteMsgToTCP(c io.Writer, m *dns.Msg) (n int, err error) {
-	mRaw, buf, err := PackBuffer(m)
+	mRaw, buf, err := pool.PackBuffer(m)
 	if err != nil {
 		return 0, err
 	}
-	defer ReleaseMsgBuf(buf)
+	defer pool.ReleaseMsgBuf(buf)
 
 	return WriteRawMsgToTCP(c, mRaw)
 }
@@ -222,5 +223,5 @@ func WriteRawMsgToTCP(c io.Writer, b []byte) (n int, err error) {
 }
 
 var (
-	tcpWriteBufPool = NewBytesBufPool(512 + 2)
+	tcpWriteBufPool = pool.NewBytesBufPool(512 + 2)
 )

dispatcher/utils/net_io_test.go → dispatcher/pkg/dnsutils/net_io_test.go

@@ -14,7 +14,8 @@
 //
 //     You should have received a copy of the GNU General Public License
 //     along with this program.  If not, see <https://www.gnu.org/licenses/>.
-package utils
+
+package dnsutils
 
 import (
 	"bytes"