diff --git a/utils/hub.go b/utils/hub.go new file mode 100644 index 0000000000..ce3f09e623 --- /dev/null +++ b/utils/hub.go @@ -0,0 +1,53 @@ +package utils + +import ( + "fmt" + "github.com/docker/beam" + "github.com/docker/beam/inmem" + "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) Wait() { + hub.tasks.Wait() +} + +func (hub *Hub) Close() error { + return hub.handlers.Close() +} diff --git a/utils/hub_test.go b/utils/hub_test.go new file mode 100644 index 0000000000..466e7c0cfc --- /dev/null +++ b/utils/hub_test.go @@ -0,0 +1,49 @@ +package utils + +import ( + "github.com/docker/beam" + "github.com/dotcloud/docker/pkg/testutils" + "testing" +) + +func TestHubSendEmpty(t *testing.T) { + hub := NewHub() + // Send to empty hub should silently drop + if r, w, err := hub.Send(&beam.Message{"hello", nil}, beam.R|beam.W); err != nil { + t.Fatal(err) + } else if r != nil { + t.Fatalf("%#v", r) + } else if w != nil { + t.Fatalf("%#v", w) + } +} + +type CountSender int + +func (s *CountSender) Send(msg *beam.Message, mode int) (beam.Receiver, beam.Sender, error) { + (*s)++ + return nil, nil, nil +} + +func TestHubSendOneHandler(t *testing.T) { + hub := NewHub() + defer hub.Close() + testutils.Timeout(t, func() { + in, _, err := hub.Send(&beam.Message{"register", nil}, beam.R) + if err != nil { + t.Fatal(err) + } + go func() { + if _, _, err := hub.Send(&beam.Message{"hello", nil}, 0); err != nil { + t.Fatal(err) + } + }() + msg, _, _, err := in.Receive(0) + if err != nil { + t.Fatal(err) + } + if msg.Name != "hello" { + t.Fatalf("%#v", msg) + } + }) +}