github.com/marinho/drone@v0.2.1-0.20140504195434-d3ba962e89a7/pkg/channel/hub.go (about)

     1  package channel
     2  
     3  import (
     4  	"sync"
     5  )
     6  
     7  // mutex to lock access to the
     8  // internal map of hubs.
     9  var mu sync.RWMutex
    10  
    11  // a map of hubs. each hub represents a different
    12  // channel that a set of users can listen on. For
    13  // example, we may have a hub to stream build output
    14  // for github.com/foo/bar or a channel to post
    15  // updates for user octocat.
    16  var hubs = map[string]*hub{}
    17  
    18  type hub struct {
    19  	// Registered connections
    20  	connections map[*connection]bool
    21  
    22  	// Inbound messages from the connections.
    23  	broadcast chan string
    24  
    25  	// Register requests from the connections.
    26  	register chan *connection
    27  
    28  	// Unregister requests from connections.
    29  	unregister chan *connection
    30  
    31  	// Buffer of sent data. This is used mostly
    32  	// for build output. A client may connect after
    33  	// the build has already started, in which case
    34  	// we need to stream them the build history.
    35  	history []string
    36  
    37  	// Send a "shutdown" signal
    38  	close chan bool
    39  
    40  	// Hub responds on this channel letting you know
    41  	// if it's active
    42  	closed chan bool
    43  
    44  	// Auto shutdown when last connection removed
    45  	autoClose bool
    46  
    47  	// Send history
    48  	sendHistory bool
    49  }
    50  
    51  func newHub(sendHistory, autoClose bool) *hub {
    52  	h := hub{
    53  		broadcast:   make(chan string),
    54  		register:    make(chan *connection),
    55  		unregister:  make(chan *connection),
    56  		connections: make(map[*connection]bool),
    57  		history:     make([]string, 0), // This should be pre-allocated, but it's not
    58  		close:       make(chan bool),
    59  		autoClose:   autoClose,
    60  		closed:      make(chan bool),
    61  		sendHistory: sendHistory,
    62  	}
    63  
    64  	return &h
    65  }
    66  
    67  func sendHistory(c *connection, history []string) {
    68  	if len(history) > 0 {
    69  		for i := range history {
    70  			c.send <- history[i]
    71  		}
    72  	}
    73  }
    74  
    75  func (h *hub) run() {
    76  	// make sure we don't bring down the application
    77  	// if somehow we encounter a nil pointer or some
    78  	// other unexpected behavior.
    79  	defer func() {
    80  		recover()
    81  	}()
    82  
    83  	for {
    84  		select {
    85  		case c := <-h.register:
    86  			h.connections[c] = true
    87  			if len(h.history) > 0 {
    88  				b := make([]string, len(h.history))
    89  				copy(b, h.history)
    90  				go sendHistory(c, b)
    91  			}
    92  		case c := <-h.unregister:
    93  			delete(h.connections, c)
    94  			close(c.send)
    95  			shutdown := h.autoClose && (len(h.connections) == 0)
    96  			if shutdown {
    97  				h.closed <- shutdown
    98  				return
    99  			}
   100  			h.closed <- shutdown
   101  		case m := <-h.broadcast:
   102  			if h.sendHistory {
   103  				h.history = append(h.history, m)
   104  			}
   105  			for c := range h.connections {
   106  				select {
   107  				case c.send <- m:
   108  					// do nothing
   109  				default:
   110  					delete(h.connections, c)
   111  					go c.ws.Close()
   112  				}
   113  			}
   114  		case <-h.close:
   115  			for c := range h.connections {
   116  				delete(h.connections, c)
   117  				close(c.send)
   118  			}
   119  			h.closed <- true
   120  			return
   121  		}
   122  
   123  	}
   124  }
   125  
   126  func (h *hub) Close() {
   127  	h.close <- true
   128  }
   129  
   130  func (h *hub) Write(p []byte) (n int, err error) {
   131  	h.broadcast <- string(p)
   132  	return len(p), nil
   133  }