working worker

object.go

@@ -29,7 +29,6 @@ type Object struct {
 
 func GetObject(s ObjectStore, id ObjectID) ([]byte, error) {
 	c := make(chan ObjectResult)
-	defer close(c)
 	ids := []ObjectID{id}
 	go func() {
 		s.AwaitObjects(ids, c, nil)
@@ -46,32 +45,92 @@ func GetObject(s ObjectStore, id ObjectID) ([]byte, error) {
 
 type MemObjectStore struct {
 	sync.RWMutex
-	db      map[ObjectID][]byte
-	prefix  uint64
-	printer uint64
+	db          map[ObjectID][]byte
+	prefix      uint64
+	printer     uint64
+	subscribers []chan *Object
 }
 
 func (mem *MemObjectStore) AwaitObjects(ids []ObjectID, c chan ObjectResult, timeout *time.Duration) {
-	if timeout != nil {
-		panic("timeout not yet implemented")
-	}
-	mem.RLock()
-	defer mem.RUnlock()
+	waitchan := make(chan *Object)
+	found := 0
+	waitingFor := make(map[ObjectID]bool)
+
+	// First, get everything we can
+	mem.Lock()
 	for _, id := range ids {
 		v, ok := mem.db[id]
 		if !ok {
-			c <- ObjectResult{Error: ErrObjectNotFound}
+			waitingFor[id] = true
 		} else {
+			found++
 			c <- ObjectResult{&Object{ID: id, Data: v}, nil}
 		}
 	}
+
+	// We were lucky, and are done.
+	if found == len(ids) {
+		mem.Unlock()
+		close(c)
+		return
+	}
+
+	// Wait for the rest
+	mem.subscribe(waitchan)
+	mem.Unlock()
+
+	var timer <-chan time.Time
+	if timeout != nil {
+		timer = time.After(*timeout)
+	} else {
+		timer = make(<-chan time.Time)
+	}
+	for found != len(ids) {
+		giveUp := false
+		select {
+		case o := <-waitchan:
+			if waitingFor[o.ID] {
+				c <- ObjectResult{o, nil}
+				waitingFor[o.ID] = false
+				found += 1
+			}
+		case <-timer:
+			giveUp = true
+		}
+		if giveUp {
+			break
+		}
+	}
+	mem.Unsubscribe(waitchan)
 	close(c)
 }
 
+func (mem *MemObjectStore) subscribe(c chan *Object) {
+	mem.subscribers = append(mem.subscribers, c)
+}
+
+func (mem *MemObjectStore) Unsubscribe(c chan *Object) {
+	mem.Lock()
+	defer mem.Unlock()
+	mem.unsubscribe(c)
+}
+
+func (mem *MemObjectStore) unsubscribe(c chan *Object) {
+	for i, s := range mem.subscribers {
+		if s == c {
+			mem.subscribers = append(mem.subscribers[:i], mem.subscribers[i+1:]...)
+			break
+		}
+	}
+}
+
 func (mem *MemObjectStore) PutObject(object *Object) error {
 	mem.Lock()
 	defer mem.Unlock()
 	mem.db[object.ID] = object.Data
+	for _, s := range mem.subscribers {
+		s <- object
+	}
 	return nil
 }