mirror of
https://github.com/fhmq/hmq.git
synced 2026-05-02 14:28:34 +00:00
Compare commits
8 Commits
master_cod
...
context
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
7ce2e3d185 | ||
|
|
684584b208 | ||
|
|
56fb4a2d54 | ||
|
|
5ed4728575 | ||
|
|
c0fea6a5ba | ||
|
|
47500910e1 | ||
|
|
0ff20b6ee2 | ||
|
|
7155667f6c |
60
README.md
60
README.md
@@ -16,43 +16,26 @@ $ go run main.go
|
|||||||
|
|
||||||
## Usage of hmq:
|
## Usage of hmq:
|
||||||
~~~
|
~~~
|
||||||
Usage of ./hmq:
|
Usage: hmq [options]
|
||||||
-w int
|
|
||||||
worker num to process message, perfer (client num)/10. (default 1024)
|
Broker Options:
|
||||||
-worker int
|
-w, --worker <number> Worker num to process message, perfer (client num)/10. (default 1024)
|
||||||
worker num to process message, perfer (client num)/10. (default 1024)
|
-p, --port <port> Use port for clients (default: 1883)
|
||||||
-h string
|
--host <host> Network host to listen on. (default "0.0.0.0")
|
||||||
Network host to listen on. (default "0.0.0.0")
|
-ws, --wsport <port> Use port for websocket monitoring
|
||||||
-host string
|
-wsp,--wspath <path> Use path for websocket monitoring
|
||||||
Network host to listen on. (default "0.0.0.0")
|
-c, --config <file> Configuration file
|
||||||
-p string
|
|
||||||
Port to listen on. (default "1883")
|
Logging Options:
|
||||||
-port string
|
-d, --debug <bool> Enable debugging output (default false)
|
||||||
Port to listen on. (default "1883")
|
-D Debug enabled
|
||||||
-c string
|
|
||||||
config file for hmq
|
Cluster Options:
|
||||||
-config string
|
-r, --router <rurl> Router who maintenance cluster info
|
||||||
config file for hmq
|
-cp, --clusterport <cluster-port> Cluster listen port for others
|
||||||
-cluster string
|
|
||||||
Cluster ip from which members can connect.
|
Common Options:
|
||||||
-cluster_listen string
|
-h, --help Show this message
|
||||||
Cluster ip from which members can connect.
|
|
||||||
-cluster_port string
|
|
||||||
Cluster port from which members can connect.
|
|
||||||
-cp string
|
|
||||||
Cluster port from which members can connect.
|
|
||||||
-r string
|
|
||||||
Router who maintenance cluster info
|
|
||||||
-router string
|
|
||||||
Router who maintenance cluster info
|
|
||||||
-ws_path string
|
|
||||||
path for ws to listen on
|
|
||||||
-ws_port string
|
|
||||||
port for ws to listen on
|
|
||||||
-wspath string
|
|
||||||
path for ws to listen on
|
|
||||||
-wsport string
|
|
||||||
port for ws to listen on
|
|
||||||
~~~
|
~~~
|
||||||
|
|
||||||
### hmq.config
|
### hmq.config
|
||||||
@@ -105,6 +88,9 @@ Usage of ./hmq:
|
|||||||
### Cluster
|
### Cluster
|
||||||
```bash
|
```bash
|
||||||
1, start router for hmq (https://github.com/fhmq/router.git)
|
1, start router for hmq (https://github.com/fhmq/router.git)
|
||||||
|
$ go get github.com/fhmq/router
|
||||||
|
$ cd $GOPATH/github.com/fhmq/router
|
||||||
|
$ go run main.go
|
||||||
2, config router in hmq.config ("router": "127.0.0.1:9888")
|
2, config router in hmq.config ("router": "127.0.0.1:9888")
|
||||||
|
|
||||||
```
|
```
|
||||||
|
|||||||
@@ -3,12 +3,10 @@
|
|||||||
package broker
|
package broker
|
||||||
|
|
||||||
import (
|
import (
|
||||||
"strings"
|
|
||||||
|
|
||||||
"github.com/fhmq/hmq/lib/acl"
|
"github.com/fhmq/hmq/lib/acl"
|
||||||
"go.uber.org/zap"
|
|
||||||
|
|
||||||
"github.com/fsnotify/fsnotify"
|
"github.com/fsnotify/fsnotify"
|
||||||
|
"go.uber.org/zap"
|
||||||
|
"strings"
|
||||||
)
|
)
|
||||||
|
|
||||||
const (
|
const (
|
||||||
|
|||||||
130
broker/broker.go
130
broker/broker.go
@@ -11,45 +11,63 @@ import (
|
|||||||
"sync/atomic"
|
"sync/atomic"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
"github.com/fhmq/hmq/lib/acl"
|
|
||||||
|
|
||||||
"github.com/eclipse/paho.mqtt.golang/packets"
|
"github.com/eclipse/paho.mqtt.golang/packets"
|
||||||
|
"github.com/fhmq/hmq/lib/acl"
|
||||||
|
"github.com/fhmq/hmq/pool"
|
||||||
"github.com/shirou/gopsutil/mem"
|
"github.com/shirou/gopsutil/mem"
|
||||||
"go.uber.org/zap"
|
"go.uber.org/zap"
|
||||||
|
|
||||||
"golang.org/x/net/websocket"
|
"golang.org/x/net/websocket"
|
||||||
|
|
||||||
"github.com/fhmq/hmq/logger"
|
|
||||||
)
|
)
|
||||||
|
|
||||||
var (
|
const (
|
||||||
log = logger.Get().Named("Broker")
|
MessagePoolNum = 1024
|
||||||
|
MessagePoolMessageNum = 1024
|
||||||
)
|
)
|
||||||
|
|
||||||
|
type Message struct {
|
||||||
|
client *client
|
||||||
|
packet packets.ControlPacket
|
||||||
|
}
|
||||||
|
|
||||||
type Broker struct {
|
type Broker struct {
|
||||||
id string
|
id string
|
||||||
cid uint64
|
cid uint64
|
||||||
mu sync.Mutex
|
mu sync.Mutex
|
||||||
config *Config
|
config *Config
|
||||||
tlsConfig *tls.Config
|
tlsConfig *tls.Config
|
||||||
AclConfig *acl.ACLConfig
|
AclConfig *acl.ACLConfig
|
||||||
clients sync.Map
|
wpool *pool.WorkerPool
|
||||||
routes sync.Map
|
clients sync.Map
|
||||||
remotes sync.Map
|
routes sync.Map
|
||||||
nodes map[string]interface{}
|
remotes sync.Map
|
||||||
sl *Sublist
|
nodes map[string]interface{}
|
||||||
rl *RetainList
|
clusterPool chan *Message
|
||||||
queues map[string]int
|
sl *Sublist
|
||||||
|
rl *RetainList
|
||||||
|
queues map[string]int
|
||||||
|
// messagePool []chan *Message
|
||||||
|
}
|
||||||
|
|
||||||
|
func newMessagePool() []chan *Message {
|
||||||
|
pool := make([]chan *Message, 0)
|
||||||
|
for i := 0; i < MessagePoolNum; i++ {
|
||||||
|
ch := make(chan *Message, MessagePoolMessageNum)
|
||||||
|
pool = append(pool, ch)
|
||||||
|
}
|
||||||
|
return pool
|
||||||
}
|
}
|
||||||
|
|
||||||
func NewBroker(config *Config) (*Broker, error) {
|
func NewBroker(config *Config) (*Broker, error) {
|
||||||
b := &Broker{
|
b := &Broker{
|
||||||
id: GenUniqueId(),
|
id: GenUniqueId(),
|
||||||
config: config,
|
config: config,
|
||||||
sl: NewSublist(),
|
wpool: pool.New(config.Worker),
|
||||||
rl: NewRetainList(),
|
sl: NewSublist(),
|
||||||
nodes: make(map[string]interface{}),
|
rl: NewRetainList(),
|
||||||
queues: make(map[string]int),
|
nodes: make(map[string]interface{}),
|
||||||
|
queues: make(map[string]int),
|
||||||
|
clusterPool: make(chan *Message),
|
||||||
|
// messagePool: newMessagePool(),
|
||||||
}
|
}
|
||||||
if b.config.TlsPort != "" {
|
if b.config.TlsPort != "" {
|
||||||
tlsconfig, err := NewTLSConfig(b.config.TlsInfo)
|
tlsconfig, err := NewTLSConfig(b.config.TlsInfo)
|
||||||
@@ -71,12 +89,26 @@ func NewBroker(config *Config) (*Broker, error) {
|
|||||||
return b, nil
|
return b, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (b *Broker) SubmitWork(msg *Message) {
|
||||||
|
if b.wpool == nil {
|
||||||
|
b.wpool = pool.New(b.config.Worker)
|
||||||
|
}
|
||||||
|
|
||||||
|
if msg.client.typ == CLUSTER {
|
||||||
|
b.clusterPool <- msg
|
||||||
|
} else {
|
||||||
|
b.wpool.Submit(func() {
|
||||||
|
ProcessMessage(msg)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
}
|
||||||
|
|
||||||
func (b *Broker) Start() {
|
func (b *Broker) Start() {
|
||||||
if b == nil {
|
if b == nil {
|
||||||
log.Error("broker is null")
|
log.Error("broker is null")
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
StartDispatcher()
|
|
||||||
|
|
||||||
//listen clinet over tcp
|
//listen clinet over tcp
|
||||||
if b.config.Port != "" {
|
if b.config.Port != "" {
|
||||||
@@ -100,6 +132,7 @@ func (b *Broker) Start() {
|
|||||||
|
|
||||||
//connect on other node in cluster
|
//connect on other node in cluster
|
||||||
if b.config.Router != "" {
|
if b.config.Router != "" {
|
||||||
|
go b.processClusterInfo()
|
||||||
b.ConnectToDiscovery()
|
b.ConnectToDiscovery()
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -124,7 +157,7 @@ func StateMonitor() {
|
|||||||
func (b *Broker) StartWebsocketListening() {
|
func (b *Broker) StartWebsocketListening() {
|
||||||
path := b.config.WsPath
|
path := b.config.WsPath
|
||||||
hp := ":" + b.config.WsPort
|
hp := ":" + b.config.WsPort
|
||||||
log.Info("Start Websocket Listening on ", zap.String("hp", hp), zap.String("path", path))
|
log.Info("Start Websocket Listener on:", zap.String("hp", hp), zap.String("path", path))
|
||||||
http.Handle(path, websocket.Handler(b.wsHandler))
|
http.Handle(path, websocket.Handler(b.wsHandler))
|
||||||
var err error
|
var err error
|
||||||
if b.config.WsTLS {
|
if b.config.WsTLS {
|
||||||
@@ -133,7 +166,7 @@ func (b *Broker) StartWebsocketListening() {
|
|||||||
err = http.ListenAndServe(hp, nil)
|
err = http.ListenAndServe(hp, nil)
|
||||||
}
|
}
|
||||||
if err != nil {
|
if err != nil {
|
||||||
log.Error("ListenAndServe: " + err.Error())
|
log.Error("ListenAndServe:" + err.Error())
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -142,7 +175,7 @@ func (b *Broker) wsHandler(ws *websocket.Conn) {
|
|||||||
// io.Copy(ws, ws)
|
// io.Copy(ws, ws)
|
||||||
atomic.AddUint64(&b.cid, 1)
|
atomic.AddUint64(&b.cid, 1)
|
||||||
ws.PayloadType = websocket.BinaryFrame
|
ws.PayloadType = websocket.BinaryFrame
|
||||||
b.handleConnection(CLIENT, ws, b.cid)
|
b.handleConnection(CLIENT, ws)
|
||||||
}
|
}
|
||||||
|
|
||||||
func (b *Broker) StartClientListening(Tls bool) {
|
func (b *Broker) StartClientListening(Tls bool) {
|
||||||
@@ -181,7 +214,7 @@ func (b *Broker) StartClientListening(Tls bool) {
|
|||||||
}
|
}
|
||||||
tmpDelay = ACCEPT_MIN_SLEEP
|
tmpDelay = ACCEPT_MIN_SLEEP
|
||||||
atomic.AddUint64(&b.cid, 1)
|
atomic.AddUint64(&b.cid, 1)
|
||||||
go b.handleConnection(CLIENT, conn, b.cid)
|
go b.handleConnection(CLIENT, conn)
|
||||||
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -225,7 +258,6 @@ func (b *Broker) StartClusterListening() {
|
|||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
var idx uint64 = 0
|
|
||||||
tmpDelay := 10 * ACCEPT_MIN_SLEEP
|
tmpDelay := 10 * ACCEPT_MIN_SLEEP
|
||||||
for {
|
for {
|
||||||
conn, err := l.Accept()
|
conn, err := l.Accept()
|
||||||
@@ -245,11 +277,11 @@ func (b *Broker) StartClusterListening() {
|
|||||||
}
|
}
|
||||||
tmpDelay = ACCEPT_MIN_SLEEP
|
tmpDelay = ACCEPT_MIN_SLEEP
|
||||||
|
|
||||||
go b.handleConnection(ROUTER, conn, idx)
|
go b.handleConnection(ROUTER, conn)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func (b *Broker) handleConnection(typ int, conn net.Conn, idx uint64) {
|
func (b *Broker) handleConnection(typ int, conn net.Conn) {
|
||||||
//process connect packet
|
//process connect packet
|
||||||
packet, err := packets.ReadPacket(conn)
|
packet, err := packets.ReadPacket(conn)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
@@ -303,39 +335,34 @@ func (b *Broker) handleConnection(typ int, conn net.Conn, idx uint64) {
|
|||||||
|
|
||||||
cid := c.info.clientID
|
cid := c.info.clientID
|
||||||
|
|
||||||
var msgPool *MessagePool
|
|
||||||
var exist bool
|
var exist bool
|
||||||
var old interface{}
|
var old interface{}
|
||||||
|
|
||||||
switch typ {
|
switch typ {
|
||||||
case CLIENT:
|
case CLIENT:
|
||||||
msgPool = MSGPool[idx%MessagePoolNum].GetPool()
|
|
||||||
c.mp = msgPool
|
|
||||||
old, exist = b.clients.Load(cid)
|
old, exist = b.clients.Load(cid)
|
||||||
if exist {
|
if exist {
|
||||||
log.Warn("client exist, close old...", zap.String("clientID", c.info.clientID))
|
log.Warn("client exist, close old...", zap.String("clientID", c.info.clientID))
|
||||||
ol, ok := old.(*client)
|
ol, ok := old.(*client)
|
||||||
if ok {
|
if ok {
|
||||||
msg := &Message{client: c, packet: DisconnectdPacket}
|
ol.Close()
|
||||||
ol.mp.queue <- msg
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
b.clients.Store(cid, c)
|
b.clients.Store(cid, c)
|
||||||
case ROUTER:
|
case ROUTER:
|
||||||
msgPool = MSGPool[(MessagePoolNum + idx)].GetPool()
|
|
||||||
c.mp = msgPool
|
|
||||||
old, exist = b.routes.Load(cid)
|
old, exist = b.routes.Load(cid)
|
||||||
if exist {
|
if exist {
|
||||||
log.Warn("router exist, close old...")
|
log.Warn("router exist, close old...")
|
||||||
ol, ok := old.(*client)
|
ol, ok := old.(*client)
|
||||||
if ok {
|
if ok {
|
||||||
msg := &Message{client: c, packet: DisconnectdPacket}
|
ol.Close()
|
||||||
ol.mp.queue <- msg
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
b.routes.Store(cid, c)
|
b.routes.Store(cid, c)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// mpool := b.messagePool[fnv1a.HashString64(cid)%MessagePoolNum]
|
||||||
|
|
||||||
c.readLoop()
|
c.readLoop()
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -383,11 +410,22 @@ func (b *Broker) ConnectToDiscovery() {
|
|||||||
c.SendConnect()
|
c.SendConnect()
|
||||||
c.SendInfo()
|
c.SendInfo()
|
||||||
|
|
||||||
c.mp = &MSGPool[(MessagePoolNum + 2)]
|
|
||||||
go c.readLoop()
|
go c.readLoop()
|
||||||
go c.StartPing()
|
go c.StartPing()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (b *Broker) processClusterInfo() {
|
||||||
|
for {
|
||||||
|
msg, ok := <-b.clusterPool
|
||||||
|
if !ok {
|
||||||
|
log.Error("read message from cluster channel error")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
ProcessMessage(msg)
|
||||||
|
}
|
||||||
|
|
||||||
|
}
|
||||||
|
|
||||||
func (b *Broker) connectRouter(id, addr string) {
|
func (b *Broker) connectRouter(id, addr string) {
|
||||||
var conn net.Conn
|
var conn net.Conn
|
||||||
var err error
|
var err error
|
||||||
@@ -446,11 +484,9 @@ func (b *Broker) connectRouter(id, addr string) {
|
|||||||
c.init()
|
c.init()
|
||||||
b.remotes.Store(cid, c)
|
b.remotes.Store(cid, c)
|
||||||
|
|
||||||
c.mp = MSGPool[(MessagePoolNum + 1)].GetPool()
|
|
||||||
|
|
||||||
c.SendConnect()
|
c.SendConnect()
|
||||||
// c.SendInfo()
|
|
||||||
|
|
||||||
|
// mpool := b.messagePool[fnv1a.HashString64(cid)%MessagePoolNum]
|
||||||
go c.readLoop()
|
go c.readLoop()
|
||||||
go c.StartPing()
|
go c.StartPing()
|
||||||
|
|
||||||
|
|||||||
108
broker/client.go
108
broker/client.go
@@ -3,7 +3,10 @@
|
|||||||
package broker
|
package broker
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"context"
|
||||||
|
"errors"
|
||||||
"net"
|
"net"
|
||||||
|
"reflect"
|
||||||
"strings"
|
"strings"
|
||||||
"sync"
|
"sync"
|
||||||
"time"
|
"time"
|
||||||
@@ -29,18 +32,18 @@ const (
|
|||||||
)
|
)
|
||||||
|
|
||||||
type client struct {
|
type client struct {
|
||||||
typ int
|
typ int
|
||||||
mu sync.Mutex
|
mu sync.Mutex
|
||||||
broker *Broker
|
broker *Broker
|
||||||
conn net.Conn
|
conn net.Conn
|
||||||
info info
|
info info
|
||||||
route route
|
route route
|
||||||
status int
|
status int
|
||||||
closed chan int
|
smu sync.RWMutex
|
||||||
smu sync.RWMutex
|
subs map[string]*subscription
|
||||||
mp *MessagePool
|
rsubs map[string]*subInfo
|
||||||
subs map[string]*subscription
|
ctx context.Context
|
||||||
rsubs map[string]*subInfo
|
cancelFunc context.CancelFunc
|
||||||
}
|
}
|
||||||
|
|
||||||
type subInfo struct {
|
type subInfo struct {
|
||||||
@@ -78,19 +81,20 @@ func (c *client) init() {
|
|||||||
c.smu.Lock()
|
c.smu.Lock()
|
||||||
defer c.smu.Unlock()
|
defer c.smu.Unlock()
|
||||||
c.status = Connected
|
c.status = Connected
|
||||||
c.closed = make(chan int, 1)
|
|
||||||
|
|
||||||
c.rsubs = make(map[string]*subInfo)
|
c.rsubs = make(map[string]*subInfo)
|
||||||
c.subs = make(map[string]*subscription, 10)
|
c.subs = make(map[string]*subscription, 10)
|
||||||
c.info.localIP = strings.Split(c.conn.LocalAddr().String(), ":")[0]
|
c.info.localIP = strings.Split(c.conn.LocalAddr().String(), ":")[0]
|
||||||
c.info.remoteIP = strings.Split(c.conn.RemoteAddr().String(), ":")[0]
|
c.info.remoteIP = strings.Split(c.conn.RemoteAddr().String(), ":")[0]
|
||||||
|
c.ctx, c.cancelFunc = context.WithCancel(context.Background())
|
||||||
}
|
}
|
||||||
|
|
||||||
func (c *client) keepAlive(ch chan int) {
|
func (c *client) keepAlive(ch chan int) {
|
||||||
defer close(ch)
|
defer close(ch)
|
||||||
|
|
||||||
|
b := c.broker
|
||||||
|
|
||||||
keepalive := time.Duration(c.info.keepalive*3/2) * time.Second
|
keepalive := time.Duration(c.info.keepalive*3/2) * time.Second
|
||||||
timer := time.NewTimer(keepalive)
|
timer := time.NewTimer(keepalive)
|
||||||
msgPool := c.mp
|
|
||||||
|
|
||||||
for {
|
for {
|
||||||
select {
|
select {
|
||||||
@@ -102,22 +106,22 @@ func (c *client) keepAlive(ch chan int) {
|
|||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
log.Error("Client exceeded timeout, disconnecting. ", zap.String("ClientID", c.info.clientID), zap.Uint16("keepalive", c.info.keepalive))
|
log.Error("Client exceeded timeout, disconnecting. ", zap.String("ClientID", c.info.clientID), zap.Uint16("keepalive", c.info.keepalive))
|
||||||
|
|
||||||
msg := &Message{client: c, packet: DisconnectdPacket}
|
msg := &Message{client: c, packet: DisconnectdPacket}
|
||||||
msgPool.queue <- msg
|
b.SubmitWork(msg)
|
||||||
|
|
||||||
timer.Stop()
|
timer.Stop()
|
||||||
return
|
return
|
||||||
case _, ok := <-c.closed:
|
case <-c.ctx.Done():
|
||||||
if !ok {
|
return
|
||||||
return
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func (c *client) readLoop() {
|
func (c *client) readLoop() {
|
||||||
nc := c.conn
|
nc := c.conn
|
||||||
msgPool := c.mp
|
b := c.broker
|
||||||
if nc == nil || msgPool == nil {
|
if nc == nil || b == nil {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -125,23 +129,28 @@ func (c *client) readLoop() {
|
|||||||
go c.keepAlive(ch)
|
go c.keepAlive(ch)
|
||||||
|
|
||||||
for {
|
for {
|
||||||
packet, err := packets.ReadPacket(nc)
|
select {
|
||||||
if err != nil {
|
case <-c.ctx.Done():
|
||||||
log.Error("read packet error: ", zap.Error(err), zap.String("ClientID", c.info.clientID))
|
return
|
||||||
break
|
default:
|
||||||
}
|
packet, err := packets.ReadPacket(nc)
|
||||||
|
if err != nil {
|
||||||
|
log.Error("read packet error: ", zap.Error(err), zap.String("ClientID", c.info.clientID))
|
||||||
|
msg := &Message{client: c, packet: DisconnectdPacket}
|
||||||
|
b.SubmitWork(msg)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
// keepalive channel
|
||||||
|
ch <- 1
|
||||||
|
|
||||||
ch <- 1
|
msg := &Message{
|
||||||
|
client: c,
|
||||||
msg := &Message{
|
packet: packet,
|
||||||
client: c,
|
}
|
||||||
packet: packet,
|
b.SubmitWork(msg)
|
||||||
}
|
}
|
||||||
msgPool.queue <- msg
|
|
||||||
}
|
}
|
||||||
msg := &Message{client: c, packet: DisconnectdPacket}
|
|
||||||
msgPool.queue <- msg
|
|
||||||
msgPool.Reduce()
|
|
||||||
}
|
}
|
||||||
|
|
||||||
func ProcessMessage(msg *Message) {
|
func ProcessMessage(msg *Message) {
|
||||||
@@ -150,10 +159,10 @@ func ProcessMessage(msg *Message) {
|
|||||||
if ca == nil {
|
if ca == nil {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
log.Debug("Recv message from client,", zap.String("ClientID", c.info.clientID))
|
|
||||||
|
log.Debug("Recv message:", zap.String("message type", reflect.TypeOf(msg.packet).String()[9:]), zap.String("ClientID", c.info.clientID))
|
||||||
switch ca.(type) {
|
switch ca.(type) {
|
||||||
case *packets.ConnackPacket:
|
case *packets.ConnackPacket:
|
||||||
|
|
||||||
case *packets.ConnectPacket:
|
case *packets.ConnectPacket:
|
||||||
case *packets.PublishPacket:
|
case *packets.PublishPacket:
|
||||||
packet := ca.(*packets.PublishPacket)
|
packet := ca.(*packets.PublishPacket)
|
||||||
@@ -237,6 +246,10 @@ func (c *client) ProcessPublishMessage(packet *packets.PublishPacket) {
|
|||||||
topic := packet.TopicName
|
topic := packet.TopicName
|
||||||
|
|
||||||
r := b.sl.Match(topic)
|
r := b.sl.Match(topic)
|
||||||
|
if r == nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
// log.Info("psubs num: ", len(r.psubs))
|
// log.Info("psubs num: ", len(r.psubs))
|
||||||
if len(r.qsubs) == 0 && len(r.psubs) == 0 {
|
if len(r.qsubs) == 0 && len(r.psubs) == 0 {
|
||||||
return
|
return
|
||||||
@@ -399,6 +412,9 @@ func (c *client) ProcessSubscribe(packet *packets.SubscribePacket) {
|
|||||||
//process retain message
|
//process retain message
|
||||||
for _, t := range topics {
|
for _, t := range topics {
|
||||||
packets := b.rl.Match(t)
|
packets := b.rl.Match(t)
|
||||||
|
if packets == nil {
|
||||||
|
continue
|
||||||
|
}
|
||||||
for _, packet := range packets {
|
for _, packet := range packets {
|
||||||
log.Info("process retain message: ", zap.Any("packet", packet), zap.String("ClientID", c.info.clientID))
|
log.Info("process retain message: ", zap.Any("packet", packet), zap.String("ClientID", c.info.clientID))
|
||||||
if packet != nil {
|
if packet != nil {
|
||||||
@@ -487,9 +503,13 @@ func (c *client) Close() {
|
|||||||
c.smu.Unlock()
|
c.smu.Unlock()
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
|
c.cancelFunc()
|
||||||
|
|
||||||
|
c.status = Disconnected
|
||||||
//wait for message complete
|
//wait for message complete
|
||||||
time.Sleep(1 * time.Second)
|
time.Sleep(1 * time.Second)
|
||||||
c.status = Disconnected
|
// c.status = Disconnected
|
||||||
|
|
||||||
if c.conn != nil {
|
if c.conn != nil {
|
||||||
c.conn.Close()
|
c.conn.Close()
|
||||||
@@ -498,8 +518,6 @@ func (c *client) Close() {
|
|||||||
|
|
||||||
c.smu.Unlock()
|
c.smu.Unlock()
|
||||||
|
|
||||||
close(c.closed)
|
|
||||||
|
|
||||||
b := c.broker
|
b := c.broker
|
||||||
subs := c.subs
|
subs := c.subs
|
||||||
if b != nil {
|
if b != nil {
|
||||||
@@ -529,9 +547,17 @@ func (c *client) Close() {
|
|||||||
}
|
}
|
||||||
|
|
||||||
func (c *client) WriterPacket(packet packets.ControlPacket) error {
|
func (c *client) WriterPacket(packet packets.ControlPacket) error {
|
||||||
|
if c.status == Disconnected {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
if packet == nil {
|
if packet == nil {
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
if c.conn == nil {
|
||||||
|
c.Close()
|
||||||
|
return errors.New("connect lost ....")
|
||||||
|
}
|
||||||
|
|
||||||
c.mu.Lock()
|
c.mu.Lock()
|
||||||
err := packet.Write(c.conn)
|
err := packet.Write(c.conn)
|
||||||
|
|||||||
@@ -9,9 +9,10 @@ import (
|
|||||||
"errors"
|
"errors"
|
||||||
"flag"
|
"flag"
|
||||||
"fmt"
|
"fmt"
|
||||||
"io/ioutil"
|
"github.com/fhmq/hmq/logger"
|
||||||
|
|
||||||
"go.uber.org/zap"
|
"go.uber.org/zap"
|
||||||
|
"io/ioutil"
|
||||||
|
"os"
|
||||||
)
|
)
|
||||||
|
|
||||||
type Config struct {
|
type Config struct {
|
||||||
@@ -28,6 +29,7 @@ type Config struct {
|
|||||||
TlsInfo TLSInfo `json:"tlsInfo"`
|
TlsInfo TLSInfo `json:"tlsInfo"`
|
||||||
Acl bool `json:"acl"`
|
Acl bool `json:"acl"`
|
||||||
AclConf string `json:"aclConf"`
|
AclConf string `json:"aclConf"`
|
||||||
|
Debug bool `json:"-"`
|
||||||
}
|
}
|
||||||
|
|
||||||
type RouteInfo struct {
|
type RouteInfo struct {
|
||||||
@@ -49,30 +51,64 @@ var DefaultConfig *Config = &Config{
|
|||||||
Acl: false,
|
Acl: false,
|
||||||
}
|
}
|
||||||
|
|
||||||
func ConfigureConfig() (*Config, error) {
|
var (
|
||||||
|
log *zap.Logger
|
||||||
|
)
|
||||||
|
|
||||||
|
func showHelp() {
|
||||||
|
fmt.Printf("%s\n", usageStr)
|
||||||
|
os.Exit(0)
|
||||||
|
}
|
||||||
|
|
||||||
|
func ConfigureConfig(args []string) (*Config, error) {
|
||||||
config := &Config{}
|
config := &Config{}
|
||||||
var (
|
var (
|
||||||
|
help bool
|
||||||
configFile string
|
configFile string
|
||||||
)
|
)
|
||||||
flag.IntVar(&config.Worker, "w", 1024, "worker num to process message, perfer (client num)/10.")
|
fs := flag.NewFlagSet("hmq-broker", flag.ExitOnError)
|
||||||
flag.IntVar(&config.Worker, "worker", 1024, "worker num to process message, perfer (client num)/10.")
|
fs.Usage = showHelp
|
||||||
flag.StringVar(&config.Port, "port", "1883", "Port to listen on.")
|
|
||||||
flag.StringVar(&config.Port, "p", "1883", "Port to listen on.")
|
fs.BoolVar(&help, "h", false, "Show this message.")
|
||||||
flag.StringVar(&config.Host, "host", "0.0.0.0", "Network host to listen on.")
|
fs.BoolVar(&help, "help", false, "Show this message.")
|
||||||
flag.StringVar(&config.Host, "h", "0.0.0.0", "Network host to listen on.")
|
fs.IntVar(&config.Worker, "w", 1024, "worker num to process message, perfer (client num)/10.")
|
||||||
flag.StringVar(&config.Cluster.Host, "cluster", "", "Cluster ip from which members can connect.")
|
fs.IntVar(&config.Worker, "worker", 1024, "worker num to process message, perfer (client num)/10.")
|
||||||
flag.StringVar(&config.Cluster.Host, "cluster_listen", "", "Cluster ip from which members can connect.")
|
fs.StringVar(&config.Port, "port", "1883", "Port to listen on.")
|
||||||
flag.StringVar(&config.Cluster.Port, "cp", "", "Cluster port from which members can connect.")
|
fs.StringVar(&config.Port, "p", "1883", "Port to listen on.")
|
||||||
flag.StringVar(&config.Cluster.Port, "cluster_port", "", "Cluster port from which members can connect.")
|
fs.StringVar(&config.Host, "host", "0.0.0.0", "Network host to listen on")
|
||||||
flag.StringVar(&config.Router, "r", "", "Router who maintenance cluster info")
|
fs.StringVar(&config.Cluster.Port, "cp", "", "Cluster port from which members can connect.")
|
||||||
flag.StringVar(&config.Router, "router", "", "Router who maintenance cluster info")
|
fs.StringVar(&config.Cluster.Port, "clusterport", "", "Cluster port from which members can connect.")
|
||||||
flag.StringVar(&config.WsPort, "wsport", "", "port for ws to listen on")
|
fs.StringVar(&config.Router, "r", "", "Router who maintenance cluster info")
|
||||||
flag.StringVar(&config.WsPort, "ws_port", "", "port for ws to listen on")
|
fs.StringVar(&config.Router, "router", "", "Router who maintenance cluster info")
|
||||||
flag.StringVar(&config.WsPath, "wspath", "", "path for ws to listen on")
|
fs.StringVar(&config.WsPort, "ws", "", "port for ws to listen on")
|
||||||
flag.StringVar(&config.WsPath, "ws_path", "", "path for ws to listen on")
|
fs.StringVar(&config.WsPort, "wsport", "", "port for ws to listen on")
|
||||||
flag.StringVar(&configFile, "config", "", "config file for hmq")
|
fs.StringVar(&config.WsPath, "wsp", "", "path for ws to listen on")
|
||||||
flag.StringVar(&configFile, "c", "", "config file for hmq")
|
fs.StringVar(&config.WsPath, "wspath", "", "path for ws to listen on")
|
||||||
flag.Parse()
|
fs.StringVar(&configFile, "config", "", "config file for hmq")
|
||||||
|
fs.StringVar(&configFile, "c", "", "config file for hmq")
|
||||||
|
fs.BoolVar(&config.Debug, "debug", false, "enable Debug logging.")
|
||||||
|
fs.BoolVar(&config.Debug, "d", false, "enable Debug logging.")
|
||||||
|
|
||||||
|
fs.Bool("D", true, "enable Debug logging.")
|
||||||
|
|
||||||
|
if err := fs.Parse(args); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
if help {
|
||||||
|
showHelp()
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
fs.Visit(func(f *flag.Flag) {
|
||||||
|
switch f.Name {
|
||||||
|
case "D":
|
||||||
|
config.Debug = true
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
logger.InitLogger(config.Debug)
|
||||||
|
log = logger.Get().Named("Broker")
|
||||||
|
|
||||||
if configFile != "" {
|
if configFile != "" {
|
||||||
tmpConfig, e := LoadConfig(configFile)
|
tmpConfig, e := LoadConfig(configFile)
|
||||||
@@ -116,8 +152,6 @@ func (config *Config) check() error {
|
|||||||
config.Worker = 1024
|
config.Worker = 1024
|
||||||
}
|
}
|
||||||
|
|
||||||
WorkNum = config.Worker
|
|
||||||
|
|
||||||
if config.Port != "" {
|
if config.Port != "" {
|
||||||
if config.Host == "" {
|
if config.Host == "" {
|
||||||
config.Host = "0.0.0.0"
|
config.Host = "0.0.0.0"
|
||||||
|
|||||||
@@ -1,44 +0,0 @@
|
|||||||
/* Copyright (c) 2018, joy.zhou <chowyu08@gmail.com>
|
|
||||||
*/
|
|
||||||
package broker
|
|
||||||
|
|
||||||
var WorkNum int
|
|
||||||
|
|
||||||
type Dispatcher struct {
|
|
||||||
WorkerPool chan chan *Message
|
|
||||||
}
|
|
||||||
|
|
||||||
func StartDispatcher() {
|
|
||||||
InitMessagePool()
|
|
||||||
dispatcher := NewDispatcher()
|
|
||||||
dispatcher.Run()
|
|
||||||
}
|
|
||||||
|
|
||||||
func (d *Dispatcher) Run() {
|
|
||||||
// starting n number of workers
|
|
||||||
for i := 0; i < WorkNum; i++ {
|
|
||||||
worker := NewWorker(d.WorkerPool)
|
|
||||||
worker.Start()
|
|
||||||
}
|
|
||||||
go d.dispatch()
|
|
||||||
}
|
|
||||||
|
|
||||||
func NewDispatcher() *Dispatcher {
|
|
||||||
pool := make(chan chan *Message, WorkNum)
|
|
||||||
return &Dispatcher{WorkerPool: pool}
|
|
||||||
}
|
|
||||||
|
|
||||||
func (d *Dispatcher) dispatch() {
|
|
||||||
for i := 0; i < (MessagePoolNum + 3); i++ {
|
|
||||||
go func(idx int) {
|
|
||||||
for {
|
|
||||||
select {
|
|
||||||
case msg := <-MSGPool[idx].queue:
|
|
||||||
msgChannel := <-d.WorkerPool
|
|
||||||
msgChannel <- msg
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}(i)
|
|
||||||
}
|
|
||||||
|
|
||||||
}
|
|
||||||
@@ -6,10 +6,9 @@ import (
|
|||||||
"fmt"
|
"fmt"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
|
simplejson "github.com/bitly/go-simplejson"
|
||||||
"github.com/eclipse/paho.mqtt.golang/packets"
|
"github.com/eclipse/paho.mqtt.golang/packets"
|
||||||
"go.uber.org/zap"
|
"go.uber.org/zap"
|
||||||
|
|
||||||
simplejson "github.com/bitly/go-simplejson"
|
|
||||||
)
|
)
|
||||||
|
|
||||||
func (c *client) SendInfo() {
|
func (c *client) SendInfo() {
|
||||||
@@ -37,10 +36,8 @@ func (c *client) StartPing() {
|
|||||||
log.Error("ping error: ", zap.Error(err))
|
log.Error("ping error: ", zap.Error(err))
|
||||||
c.Close()
|
c.Close()
|
||||||
}
|
}
|
||||||
case _, ok := <-c.closed:
|
case <-c.ctx.Done():
|
||||||
if !ok {
|
return
|
||||||
return
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1,64 +0,0 @@
|
|||||||
/* Copyright (c) 2018, joy.zhou <chowyu08@gmail.com>
|
|
||||||
*/
|
|
||||||
package broker
|
|
||||||
|
|
||||||
import (
|
|
||||||
"sync"
|
|
||||||
|
|
||||||
"github.com/eclipse/paho.mqtt.golang/packets"
|
|
||||||
)
|
|
||||||
|
|
||||||
const (
|
|
||||||
MaxUser = 1024 * 1024
|
|
||||||
MessagePoolNum = 1024
|
|
||||||
MessagePoolUser = MaxUser / MessagePoolNum
|
|
||||||
MessagePoolMessageNum = MaxUser / MessagePoolNum * 4
|
|
||||||
)
|
|
||||||
|
|
||||||
type Message struct {
|
|
||||||
client *client
|
|
||||||
packet packets.ControlPacket
|
|
||||||
}
|
|
||||||
|
|
||||||
var (
|
|
||||||
MSGPool []MessagePool
|
|
||||||
)
|
|
||||||
|
|
||||||
type MessagePool struct {
|
|
||||||
l sync.Mutex
|
|
||||||
maxuser int
|
|
||||||
user int
|
|
||||||
queue chan *Message
|
|
||||||
}
|
|
||||||
|
|
||||||
func InitMessagePool() {
|
|
||||||
MSGPool = make([]MessagePool, (MessagePoolNum + 3))
|
|
||||||
for i := 0; i < (MessagePoolNum + 3); i++ {
|
|
||||||
MSGPool[i].Init(MessagePoolUser, MessagePoolMessageNum)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
func (p *MessagePool) Init(num int, maxusernum int) {
|
|
||||||
p.maxuser = maxusernum
|
|
||||||
p.queue = make(chan *Message, num)
|
|
||||||
}
|
|
||||||
|
|
||||||
func (p *MessagePool) GetPool() *MessagePool {
|
|
||||||
p.l.Lock()
|
|
||||||
if p.user+1 < p.maxuser {
|
|
||||||
p.user += 1
|
|
||||||
p.l.Unlock()
|
|
||||||
return p
|
|
||||||
} else {
|
|
||||||
p.l.Unlock()
|
|
||||||
return nil
|
|
||||||
}
|
|
||||||
|
|
||||||
}
|
|
||||||
|
|
||||||
func (p *MessagePool) Reduce() {
|
|
||||||
p.l.Lock()
|
|
||||||
p.user -= 1
|
|
||||||
p.l.Unlock()
|
|
||||||
|
|
||||||
}
|
|
||||||
@@ -1,9 +1,8 @@
|
|||||||
package broker
|
package broker
|
||||||
|
|
||||||
import (
|
import (
|
||||||
"sync"
|
|
||||||
|
|
||||||
"github.com/eclipse/paho.mqtt.golang/packets"
|
"github.com/eclipse/paho.mqtt.golang/packets"
|
||||||
|
"sync"
|
||||||
)
|
)
|
||||||
|
|
||||||
type RetainList struct {
|
type RetainList struct {
|
||||||
|
|||||||
@@ -4,9 +4,8 @@ package broker
|
|||||||
|
|
||||||
import (
|
import (
|
||||||
"errors"
|
"errors"
|
||||||
"sync"
|
|
||||||
|
|
||||||
"go.uber.org/zap"
|
"go.uber.org/zap"
|
||||||
|
"sync"
|
||||||
)
|
)
|
||||||
|
|
||||||
// A result structure better optimized for queue subs.
|
// A result structure better optimized for queue subs.
|
||||||
|
|||||||
24
broker/usage.go
Normal file
24
broker/usage.go
Normal file
@@ -0,0 +1,24 @@
|
|||||||
|
package broker
|
||||||
|
|
||||||
|
var usageStr = `
|
||||||
|
Usage: hmq [options]
|
||||||
|
|
||||||
|
Broker Options:
|
||||||
|
-w, --worker <number> Worker num to process message, perfer (client num)/10. (default 1024)
|
||||||
|
-p, --port <port> Use port for clients (default: 1883)
|
||||||
|
--host <host> Network host to listen on. (default "0.0.0.0")
|
||||||
|
-ws, --wsport <port> Use port for websocket monitoring
|
||||||
|
-wsp,--wspath <path> Use path for websocket monitoring
|
||||||
|
-c, --config <file> Configuration file
|
||||||
|
|
||||||
|
Logging Options:
|
||||||
|
-d, --debug <bool> Enable debugging output (default false)
|
||||||
|
-D Debug and trace
|
||||||
|
|
||||||
|
Cluster Options:
|
||||||
|
-r, --router <rurl> Router who maintenance cluster info
|
||||||
|
-cp, --clusterport <cluster-port> Cluster listen port for others
|
||||||
|
|
||||||
|
Common Options:
|
||||||
|
-h, --help Show this message
|
||||||
|
`
|
||||||
@@ -1,39 +0,0 @@
|
|||||||
/* Copyright (c) 2018, joy.zhou <chowyu08@gmail.com>
|
|
||||||
*/
|
|
||||||
package broker
|
|
||||||
|
|
||||||
type Worker struct {
|
|
||||||
WorkerPool chan chan *Message
|
|
||||||
MsgChannel chan *Message
|
|
||||||
quit chan bool
|
|
||||||
}
|
|
||||||
|
|
||||||
func NewWorker(workerPool chan chan *Message) Worker {
|
|
||||||
return Worker{
|
|
||||||
WorkerPool: workerPool,
|
|
||||||
MsgChannel: make(chan *Message),
|
|
||||||
quit: make(chan bool)}
|
|
||||||
}
|
|
||||||
|
|
||||||
func (w Worker) Start() {
|
|
||||||
go func() {
|
|
||||||
for {
|
|
||||||
// register the current worker into the worker queue.
|
|
||||||
w.WorkerPool <- w.MsgChannel
|
|
||||||
select {
|
|
||||||
case msg := <-w.MsgChannel:
|
|
||||||
// we have received a work request.
|
|
||||||
ProcessMessage(msg)
|
|
||||||
case <-w.quit:
|
|
||||||
return
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}()
|
|
||||||
}
|
|
||||||
|
|
||||||
// Stop signals the worker to stop listening for work requests.
|
|
||||||
func (w Worker) Stop() {
|
|
||||||
go func() {
|
|
||||||
w.quit <- true
|
|
||||||
}()
|
|
||||||
}
|
|
||||||
@@ -9,7 +9,6 @@ import (
|
|||||||
|
|
||||||
var (
|
var (
|
||||||
// env can be setup at build time with Go Linker. Value could be prod or whatever else for dev env
|
// env can be setup at build time with Go Linker. Value could be prod or whatever else for dev env
|
||||||
env string
|
|
||||||
instance *zap.Logger
|
instance *zap.Logger
|
||||||
logCfg zap.Config
|
logCfg zap.Config
|
||||||
)
|
)
|
||||||
@@ -28,13 +27,13 @@ func NewProdLogger() (*zap.Logger, error) {
|
|||||||
return logCfg.Build()
|
return logCfg.Build()
|
||||||
}
|
}
|
||||||
|
|
||||||
func init() {
|
func InitLogger(debug bool) {
|
||||||
var err error
|
var err error
|
||||||
var log *zap.Logger
|
var log *zap.Logger
|
||||||
if env == "prod" {
|
if debug {
|
||||||
log, err = NewProdLogger()
|
|
||||||
} else {
|
|
||||||
log, err = NewDevLogger()
|
log, err = NewDevLogger()
|
||||||
|
} else {
|
||||||
|
log, err = NewProdLogger()
|
||||||
}
|
}
|
||||||
if err != nil {
|
if err != nil {
|
||||||
panic("Unable to create a logger.")
|
panic("Unable to create a logger.")
|
||||||
|
|||||||
18
main.go
18
main.go
@@ -7,36 +7,30 @@ copyright notice and this permission notice appear in all copies.
|
|||||||
package main
|
package main
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"fmt"
|
||||||
|
"github.com/fhmq/hmq/broker"
|
||||||
"os"
|
"os"
|
||||||
"os/signal"
|
"os/signal"
|
||||||
"runtime"
|
"runtime"
|
||||||
|
|
||||||
"github.com/fhmq/hmq/broker"
|
|
||||||
"github.com/fhmq/hmq/logger"
|
|
||||||
"go.uber.org/zap"
|
|
||||||
)
|
|
||||||
|
|
||||||
var (
|
|
||||||
log = logger.Get().Named("Main")
|
|
||||||
)
|
)
|
||||||
|
|
||||||
func main() {
|
func main() {
|
||||||
runtime.GOMAXPROCS(runtime.NumCPU())
|
runtime.GOMAXPROCS(runtime.NumCPU())
|
||||||
config, err := broker.ConfigureConfig()
|
config, err := broker.ConfigureConfig(os.Args[1:])
|
||||||
if err != nil {
|
if err != nil {
|
||||||
log.Error("configure broker config error: ", zap.Error(err))
|
fmt.Println("configure broker config error: ", err)
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
b, err := broker.NewBroker(config)
|
b, err := broker.NewBroker(config)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
log.Error("New Broker error: ", zap.Error(err))
|
fmt.Println("New Broker error: ", err)
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
b.Start()
|
b.Start()
|
||||||
|
|
||||||
s := waitForSignal()
|
s := waitForSignal()
|
||||||
log.Info("signal received, broker closed.", zap.Any("signal", s))
|
fmt.Println("signal received, broker closed.", s)
|
||||||
}
|
}
|
||||||
|
|
||||||
func waitForSignal() os.Signal {
|
func waitForSignal() os.Signal {
|
||||||
|
|||||||
166
pool/pool.go
Normal file
166
pool/pool.go
Normal file
@@ -0,0 +1,166 @@
|
|||||||
|
package pool
|
||||||
|
|
||||||
|
import "time"
|
||||||
|
|
||||||
|
const (
|
||||||
|
// This value is the size of the queue that workers register their
|
||||||
|
// availability to the dispatcher. There may be hundreds of workers, but
|
||||||
|
// only a small channel is needed to register some of the workers.
|
||||||
|
readyQueueSize = 16
|
||||||
|
|
||||||
|
// If worker pool receives no new work for this period of time, then stop
|
||||||
|
// a worker goroutine.
|
||||||
|
idleTimeoutSec = 5
|
||||||
|
)
|
||||||
|
|
||||||
|
type WorkerPool struct {
|
||||||
|
maxWorkers int
|
||||||
|
timeout time.Duration
|
||||||
|
taskQueue chan func()
|
||||||
|
readyWorkers chan chan func()
|
||||||
|
stoppedChan chan struct{}
|
||||||
|
}
|
||||||
|
|
||||||
|
func New(maxWorkers int) *WorkerPool {
|
||||||
|
// There must be at least one worker.
|
||||||
|
if maxWorkers < 1 {
|
||||||
|
maxWorkers = 1
|
||||||
|
}
|
||||||
|
|
||||||
|
// taskQueue is unbuffered since items are always removed immediately.
|
||||||
|
pool := &WorkerPool{
|
||||||
|
taskQueue: make(chan func()),
|
||||||
|
maxWorkers: maxWorkers,
|
||||||
|
readyWorkers: make(chan chan func(), readyQueueSize),
|
||||||
|
timeout: time.Second * idleTimeoutSec,
|
||||||
|
stoppedChan: make(chan struct{}),
|
||||||
|
}
|
||||||
|
|
||||||
|
// Start the task dispatcher.
|
||||||
|
go pool.dispatch()
|
||||||
|
|
||||||
|
return pool
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *WorkerPool) Stop() {
|
||||||
|
if p.Stopped() {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
close(p.taskQueue)
|
||||||
|
<-p.stoppedChan
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *WorkerPool) Stopped() bool {
|
||||||
|
select {
|
||||||
|
case <-p.stoppedChan:
|
||||||
|
return true
|
||||||
|
default:
|
||||||
|
}
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *WorkerPool) Submit(task func()) {
|
||||||
|
if task != nil {
|
||||||
|
p.taskQueue <- task
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *WorkerPool) SubmitWait(task func()) {
|
||||||
|
if task == nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
doneChan := make(chan struct{})
|
||||||
|
p.taskQueue <- func() {
|
||||||
|
task()
|
||||||
|
close(doneChan)
|
||||||
|
}
|
||||||
|
<-doneChan
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *WorkerPool) dispatch() {
|
||||||
|
defer close(p.stoppedChan)
|
||||||
|
timeout := time.NewTimer(p.timeout)
|
||||||
|
var workerCount int
|
||||||
|
var task func()
|
||||||
|
var ok bool
|
||||||
|
var workerTaskChan chan func()
|
||||||
|
startReady := make(chan chan func())
|
||||||
|
Loop:
|
||||||
|
for {
|
||||||
|
timeout.Reset(p.timeout)
|
||||||
|
select {
|
||||||
|
case task, ok = <-p.taskQueue:
|
||||||
|
if !ok {
|
||||||
|
break Loop
|
||||||
|
}
|
||||||
|
// Got a task to do.
|
||||||
|
select {
|
||||||
|
case workerTaskChan = <-p.readyWorkers:
|
||||||
|
// A worker is ready, so give task to worker.
|
||||||
|
workerTaskChan <- task
|
||||||
|
default:
|
||||||
|
// No workers ready.
|
||||||
|
// Create a new worker, if not at max.
|
||||||
|
if workerCount < p.maxWorkers {
|
||||||
|
workerCount++
|
||||||
|
go func(t func()) {
|
||||||
|
startWorker(startReady, p.readyWorkers)
|
||||||
|
// Submit the task when the new worker.
|
||||||
|
taskChan := <-startReady
|
||||||
|
taskChan <- t
|
||||||
|
}(task)
|
||||||
|
} else {
|
||||||
|
// Start a goroutine to submit the task when an existing
|
||||||
|
// worker is ready.
|
||||||
|
go func(t func()) {
|
||||||
|
taskChan := <-p.readyWorkers
|
||||||
|
taskChan <- t
|
||||||
|
}(task)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
case <-timeout.C:
|
||||||
|
// Timed out waiting for work to arrive. Kill a ready worker.
|
||||||
|
if workerCount > 0 {
|
||||||
|
select {
|
||||||
|
case workerTaskChan = <-p.readyWorkers:
|
||||||
|
// A worker is ready, so kill.
|
||||||
|
close(workerTaskChan)
|
||||||
|
workerCount--
|
||||||
|
default:
|
||||||
|
// No work, but no ready workers. All workers are busy.
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Stop all remaining workers as they become ready.
|
||||||
|
for workerCount > 0 {
|
||||||
|
workerTaskChan = <-p.readyWorkers
|
||||||
|
close(workerTaskChan)
|
||||||
|
workerCount--
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func startWorker(startReady, readyWorkers chan chan func()) {
|
||||||
|
go func() {
|
||||||
|
taskChan := make(chan func())
|
||||||
|
var task func()
|
||||||
|
var ok bool
|
||||||
|
// Register availability on starReady channel.
|
||||||
|
startReady <- taskChan
|
||||||
|
for {
|
||||||
|
// Read task from dispatcher.
|
||||||
|
task, ok = <-taskChan
|
||||||
|
if !ok {
|
||||||
|
// Dispatcher has told worker to stop.
|
||||||
|
break
|
||||||
|
}
|
||||||
|
|
||||||
|
// Execute the task.
|
||||||
|
task()
|
||||||
|
|
||||||
|
// Register availability on readyWorkers channel.
|
||||||
|
readyWorkers <- taskChan
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user