nnet/options.go

package nnet

import (
	"git.noahlan.cn/noahlan/nnet/connection"
	"git.noahlan.cn/noahlan/nnet/lifetime"
	"git.noahlan.cn/noahlan/nnet/packet"
	rt "git.noahlan.cn/noahlan/nnet/router"
	"git.noahlan.cn/noahlan/nnet/serialize"
	"git.noahlan.cn/noahlan/ntools-go/core/pool"
	"github.com/panjf2000/ants/v2"
	"time"
)

type (
	// RunOption defines the method to customize an Engine.
	RunOption func(ngin *Engine)
)

// Pipeline returns inner pipeline
func (ngin *Engine) Pipeline() connection.Pipeline {
	return ngin.pipeline
}

// Lifetime returns lifetime interface.
func (ngin *Engine) Lifetime() lifetime.Lifetime {
	return ngin.lifetime
}

// ConnManager returns connection manager
func (ngin *Engine) ConnManager() *connection.Manager {
	return ngin.connManager
}

//////////////////////// Options

func WithMiddleware(middleware ...rt.Middleware) RunOption {
	return func(ngin *Engine) {
		ngin.Use(middleware...)
	}
}

// WithRouter 设置消息路由
func WithRouter(router rt.Router) RunOption {
	return func(ngin *Engine) {
		ngin.router = router
	}
}

// WithNotFoundHandler returns a RunOption with not found handler set to given handler.
func WithNotFoundHandler(handler rt.Handler) RunOption {
	return func(ngin *Engine) {
		ngin.router.SetNotFoundHandler(rt.NotFoundHandler(handler))
	}
}

// WithTimerPrecision 设置Timer精度，需在 Start 或 Dial 之前执行
// 注：精度需大于1ms, 并且不能在运行时更改
// 默认精度是 time.Second
func WithTimerPrecision(precision time.Duration) RunOption {
	if precision < time.Millisecond {
		panic("time precision can not less than a Millisecond")
	}
	return func(ngin *Engine) {
		ngin.TaskTimerPrecision = precision
	}
}

// WithPackerBuilder 设置 消息的封包/解包构造器
func WithPackerBuilder(fn packet.PackerBuilder) RunOption {
	return func(ngin *Engine) {
		ngin.packerBuilder = fn
	}
}

// WithSerializer 设置消息的 序列化/反序列化 方式
func WithSerializer(s serialize.Serializer) RunOption {
	return func(ngin *Engine) {
		ngin.serializer = s
	}
}

// WithPool 设置使用自定义的工作池
func WithPool(pl *ants.Pool) RunOption {
	return func(ngin *Engine) {
		ngin.goPool = pl
	}
}

// WithPoolCfg 设置工作池配置
func WithPoolCfg(cfg pool.Config) RunOption {
	return func(ngin *Engine) {
		ngin.goPool, _ = ants.NewPool(cfg.PoolSize, ants.WithOptions(cfg.Options()))
	}
}

//////////////////// Pipeline

// WithPipeline 使用自定义 pipeline
func WithPipeline(pipeline connection.Pipeline) RunOption {
	return func(ngin *Engine) {
		ngin.pipeline = pipeline
	}
}

type PipelineOption func(opts connection.Pipeline)

// WithPipelineOpt 使用默认Pipeline并设置其配置
func WithPipelineOpt(opts ...func(connection.Pipeline)) RunOption {
	return func(ngin *Engine) {
		for _, opt := range opts {
			opt(ngin.pipeline)
		}
	}
}