mirror of https://github.com/docker/docs.git
129 lines
2.8 KiB
Go
129 lines
2.8 KiB
Go
package utils
|
|
|
|
import (
|
|
"fmt"
|
|
"github.com/docker/libswarm/beam"
|
|
"github.com/docker/libswarm/beam/inmem"
|
|
"io"
|
|
"sync"
|
|
)
|
|
|
|
// Hub passes messages to dynamically registered handlers.
|
|
type Hub struct {
|
|
handlers *StackSender
|
|
tasks sync.WaitGroup
|
|
}
|
|
|
|
func NewHub() *Hub {
|
|
return &Hub{
|
|
handlers: NewStackSender(),
|
|
}
|
|
}
|
|
|
|
func (hub *Hub) Send(msg *beam.Message, mode int) (beam.Receiver, beam.Sender, error) {
|
|
if msg.Name == "register" {
|
|
if mode&beam.R == 0 {
|
|
return nil, nil, fmt.Errorf("register: no return channel")
|
|
}
|
|
fmt.Printf("[hub] received %v\n", msg)
|
|
hYoutr, hYoutw := inmem.Pipe()
|
|
hYinr, hYinw := inmem.Pipe()
|
|
// Register the new handler on top of the others,
|
|
// and get a reference to the previous stack of handlers.
|
|
prevHandlers := hub.handlers.Add(hYinw)
|
|
// Pass requests from the new handler to the previous chain of handlers
|
|
// hYout -> hXin
|
|
hub.tasks.Add(1)
|
|
go func() {
|
|
defer hub.tasks.Done()
|
|
Copy(prevHandlers, hYoutr)
|
|
hYoutr.Close()
|
|
}()
|
|
return hYinr, hYoutw, nil
|
|
}
|
|
fmt.Printf("sending %#v to %d handlers\n", msg, hub.handlers.Len())
|
|
return hub.handlers.Send(msg, mode)
|
|
}
|
|
|
|
func (hub *Hub) Register(dst beam.Sender) error {
|
|
in, _, err := hub.Send(&beam.Message{Name: "register"}, beam.R)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
go Copy(dst, in)
|
|
return nil
|
|
}
|
|
|
|
func (hub *Hub) RegisterTask(h func(beam.Receiver, beam.Sender) error) error {
|
|
in, out, err := hub.Send(&beam.Message{Name: "register"}, beam.R|beam.W)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
go func() {
|
|
h(in, out)
|
|
out.Close()
|
|
}()
|
|
return nil
|
|
}
|
|
|
|
type Handler func(msg *beam.Message, in beam.Receiver, out beam.Sender, next beam.Sender) (pass bool, err error)
|
|
|
|
func (hub *Hub) RegisterName(name string, h Handler) error {
|
|
return hub.RegisterTask(func(in beam.Receiver, out beam.Sender) error {
|
|
var tasks sync.WaitGroup
|
|
copyTask := func(dst beam.Sender, src beam.Receiver) {
|
|
tasks.Add(1)
|
|
go func() {
|
|
defer tasks.Done()
|
|
if dst == nil {
|
|
return
|
|
}
|
|
defer dst.Close()
|
|
if src == nil {
|
|
return
|
|
}
|
|
Copy(dst, src)
|
|
}()
|
|
}
|
|
for {
|
|
msg, msgin, msgout, err := in.Receive(beam.R | beam.W)
|
|
if err == io.EOF {
|
|
break
|
|
}
|
|
if err != nil {
|
|
return err
|
|
}
|
|
var pass = true
|
|
if msg.Name == name || name == "" {
|
|
pass, err = h(msg, msgin, msgout, out)
|
|
if err != nil {
|
|
if _, _, err := msgout.Send(&beam.Message{Name: "error", Args: []string{err.Error()}}, 0); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
}
|
|
if pass {
|
|
nextin, nextout, err := out.Send(msg, beam.R|beam.W)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
copyTask(nextout, msgin)
|
|
copyTask(msgout, nextin)
|
|
} else {
|
|
if msgout != nil {
|
|
msgout.Close()
|
|
}
|
|
}
|
|
}
|
|
return nil
|
|
})
|
|
}
|
|
|
|
func (hub *Hub) Wait() {
|
|
hub.tasks.Wait()
|
|
}
|
|
|
|
func (hub *Hub) Close() error {
|
|
return hub.handlers.Close()
|
|
}
|