parent
5378603349
commit
66a6f75324
@ -0,0 +1,179 @@ |
||||
package app |
||||
|
||||
import ( |
||||
"bytes" |
||||
"context" |
||||
"encoding/binary" |
||||
"fmt" |
||||
"io" |
||||
"os" |
||||
"sync" |
||||
) |
||||
|
||||
// Callback 订阅回调
|
||||
type Callback func(payload []byte) |
||||
|
||||
// Runnable 子协程运行
|
||||
type Runnable func(ctx context.Context) |
||||
|
||||
var wg sync.WaitGroup |
||||
var resChan chan []byte |
||||
var doneChan chan struct{} |
||||
var channels = make(map[string][]Callback) |
||||
var runtime = context.Background() |
||||
|
||||
var stats struct { |
||||
send int |
||||
respond int |
||||
receive int |
||||
} |
||||
|
||||
func Async(callback func()) { |
||||
wg.Add(1) |
||||
go func() { |
||||
defer wg.Done() |
||||
callback() |
||||
}() |
||||
} |
||||
|
||||
func process(buf []byte) (response []byte, err error) { |
||||
stats.receive++ |
||||
var n int |
||||
var channel []byte |
||||
n, channel, err = decode(buf) |
||||
if err != nil { |
||||
return |
||||
} |
||||
var payload []byte |
||||
_, payload, err = decode(buf[n:]) |
||||
if err != nil { |
||||
return |
||||
} |
||||
subscribers, ok := channels[string(channel)] |
||||
if !ok { |
||||
err = fmt.Errorf("no subscribers for channel %s", string(channel)) |
||||
return |
||||
} |
||||
for _, subscriber := range subscribers[:] { |
||||
// fixme: 是否需要使用 payload 的副本?
|
||||
recovery(payload, subscriber) |
||||
} |
||||
if response != nil { |
||||
stats.respond++ |
||||
} |
||||
return |
||||
} |
||||
|
||||
func recovery(payload []byte, callback Callback) { |
||||
defer func() { |
||||
err := recover() |
||||
if err != nil { |
||||
fmt.Fprintf(os.Stderr, "%v", err) |
||||
} |
||||
}() |
||||
callback(payload) |
||||
} |
||||
|
||||
func Sub(channel string, callback Callback) { |
||||
subscribers, ok := channels[channel] |
||||
if !ok { |
||||
subscribers = make([]Callback, 0) |
||||
} |
||||
subscribers = append(subscribers, callback) |
||||
channels[channel] = subscribers |
||||
} |
||||
|
||||
// 定义写数据的格式
|
||||
func encode(data []byte) []byte { |
||||
// 4字节头部 + 可变体的长度
|
||||
buf := make([]byte, 4+len(data)) |
||||
// 写入头部,记录数据长度
|
||||
binary.BigEndian.PutUint32(buf[:4], uint32(len(data))) |
||||
// 将整个数据,放到4后边
|
||||
copy(buf[4:], data) |
||||
// 返回结果
|
||||
return buf |
||||
} |
||||
|
||||
func decode(buf []byte) (int, []byte, error) { |
||||
r := bytes.NewBuffer(buf) |
||||
// 读取头部记录的长度
|
||||
header := make([]byte, 4) |
||||
// 按长度读取消息
|
||||
_, err := io.ReadFull(r, header) |
||||
if err != nil { |
||||
return 0, nil, err |
||||
} |
||||
// 读取数据
|
||||
dataLen := binary.BigEndian.Uint32(header) |
||||
data := make([]byte, dataLen) |
||||
_, err = io.ReadFull(r, data) |
||||
if err != nil { |
||||
return 0, nil, err |
||||
} |
||||
return 4 + int(dataLen), data, nil |
||||
} |
||||
|
||||
func Pub(channel string, payload []byte) { |
||||
// fixme: 在时间和空间上,是否有比 append 性能更好的操作
|
||||
buf := append(encode([]byte(channel)), encode(payload)...) |
||||
wg.Add(1) |
||||
resChan <- buf |
||||
} |
||||
|
||||
func Run(runner Runnable) { |
||||
Async(func() { |
||||
runner(runtime) |
||||
}) |
||||
} |
||||
|
||||
func Loop() { |
||||
wg.Add(1) |
||||
first := true |
||||
|
||||
var cancel context.CancelFunc |
||||
runtime, cancel = context.WithCancel(context.Background()) |
||||
|
||||
go func() { |
||||
wg.Wait() |
||||
cancel() |
||||
doneChan <- struct{}{} |
||||
}() |
||||
|
||||
for { |
||||
select { |
||||
case msg := <-resChan: |
||||
process(msg) |
||||
wg.Done() // Corresponds to the wg.Add(1) in Pub().
|
||||
case <-doneChan: |
||||
// All goroutines have completed. Now we can exit main().
|
||||
checkChanEmpty() |
||||
return |
||||
} |
||||
|
||||
// We don't want to exit until we've received at least one message.
|
||||
// This is so the program doesn't exit after sending the "start"
|
||||
// message.
|
||||
if first { |
||||
wg.Done() |
||||
first = false |
||||
} |
||||
} |
||||
|
||||
// todo 清理工作
|
||||
} |
||||
|
||||
func checkChanEmpty() { |
||||
// We've received a done event. As a sanity check, make sure that resChan is
|
||||
// empty.
|
||||
select { |
||||
case _, ok := <-resChan: |
||||
if ok { |
||||
panic("Read a message from resChan after doneChan closed.") |
||||
} else { |
||||
panic("resChan closed. Unexpected.") |
||||
} |
||||
default: |
||||
// No value ready, moving on.
|
||||
} |
||||
} |
Loading…
Reference in new issue