github.com/slspeek/camlistore_namedsearch@v0.0.0-20140519202248-ed6f70f7721a/cmd/camput/chanworker.go (about)

     1  /*
     2  Copyright 2012 Google Inc.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8       http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  // Worker pools of functions processing channel input.
    18  // TODO(brafitz): move this to a common library, making it operate on interface{} instead?
    19  
    20  package main
    21  
    22  import (
    23  	"container/list"
    24  	"sync"
    25  )
    26  
    27  type nodeWorker struct {
    28  	c chan *node
    29  
    30  	donec chan bool
    31  	workc chan *node
    32  	fn    func(n *node, ok bool)
    33  	buf   *list.List
    34  }
    35  
    36  // NewNodeWorker starts nWorkers goroutines running fn on incoming
    37  // nodes sent on the returned channel.  fn may block; writes to the
    38  // channel will buffer.
    39  // If nWorkers is negative, a new goroutine running fn is called for each
    40  // item sent on the returned channel.
    41  // When the returned channel is closed, fn is called with (nil, false)
    42  // after all other calls to fn have completed.
    43  func NewNodeWorker(nWorkers int, fn func(n *node, ok bool)) chan<- *node {
    44  	if nWorkers == 0 {
    45  		panic("invalid nWorkers valid of 0")
    46  	}
    47  	retc := make(chan *node, buffered)
    48  	if nWorkers < 0 {
    49  		// Unbounded number of workers.
    50  		go func() {
    51  			var wg sync.WaitGroup
    52  			for w := range retc {
    53  				wg.Add(1)
    54  				go func(w *node) {
    55  					fn(w, true)
    56  					wg.Done()
    57  				}(w)
    58  			}
    59  			wg.Wait()
    60  			fn(nil, false)
    61  		}()
    62  		return retc
    63  	}
    64  	w := &nodeWorker{
    65  		c:     retc,
    66  		workc: make(chan *node, buffered),
    67  		donec: make(chan bool), // when workers finish
    68  		fn:    fn,
    69  		buf:   list.New(),
    70  	}
    71  	go w.pump()
    72  	for i := 0; i < nWorkers; i++ {
    73  		go w.work()
    74  	}
    75  	go func() {
    76  		for i := 0; i < nWorkers; i++ {
    77  			<-w.donec
    78  		}
    79  		fn(nil, false) // final sentinel
    80  	}()
    81  	return retc
    82  }
    83  
    84  func (w *nodeWorker) pump() {
    85  	inc := w.c
    86  	for inc != nil || w.buf.Len() > 0 {
    87  		outc := w.workc
    88  		var frontNode *node
    89  		if e := w.buf.Front(); e != nil {
    90  			frontNode = e.Value.(*node)
    91  		} else {
    92  			outc = nil
    93  		}
    94  		select {
    95  		case outc <- frontNode:
    96  			w.buf.Remove(w.buf.Front())
    97  		case n, ok := <-inc:
    98  			if !ok {
    99  				inc = nil
   100  				continue
   101  			}
   102  			w.buf.PushBack(n)
   103  		}
   104  	}
   105  	close(w.workc)
   106  }
   107  
   108  func (w *nodeWorker) work() {
   109  	for {
   110  		select {
   111  		case n, ok := <-w.workc:
   112  			if !ok {
   113  				w.donec <- true
   114  				return
   115  			}
   116  			w.fn(n, true)
   117  		}
   118  	}
   119  }