github.com/gdamore/mangos@v1.4.0/protocol/push/push.go (about)

     1  // Copyright 2018 The Mangos Authors
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use file except in compliance with the License.
     5  // You may obtain a copy of the license at
     6  //
     7  //    http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  // Package push implements the PUSH protocol, which is the write side of
    16  // the pipeline pattern.  (PULL is the reader.)
    17  package push
    18  
    19  import (
    20  	"sync"
    21  	"time"
    22  
    23  	"nanomsg.org/go-mangos"
    24  )
    25  
    26  type push struct {
    27  	sock mangos.ProtocolSocket
    28  	raw  bool
    29  	w    mangos.Waiter
    30  	eps  map[uint32]*pushEp
    31  	sync.Mutex
    32  }
    33  
    34  type pushEp struct {
    35  	ep mangos.Endpoint
    36  	cq chan struct{}
    37  }
    38  
    39  func (x *push) Init(sock mangos.ProtocolSocket) {
    40  	x.sock = sock
    41  	x.w.Init()
    42  	x.sock.SetRecvError(mangos.ErrProtoOp)
    43  	x.eps = make(map[uint32]*pushEp)
    44  }
    45  
    46  func (x *push) Shutdown(expire time.Time) {
    47  	x.w.WaitAbsTimeout(expire)
    48  }
    49  
    50  func (x *push) sender(ep *pushEp) {
    51  	defer x.w.Done()
    52  	sq := x.sock.SendChannel()
    53  	cq := x.sock.CloseChannel()
    54  
    55  	for {
    56  		select {
    57  		case <-cq:
    58  			return
    59  		case <-ep.cq:
    60  			return
    61  		case m := <-sq:
    62  			if m == nil {
    63  				sq = x.sock.SendChannel()
    64  				continue
    65  			}
    66  			if ep.ep.SendMsg(m) != nil {
    67  				m.Free()
    68  				return
    69  			}
    70  		}
    71  	}
    72  }
    73  
    74  func (*push) Number() uint16 {
    75  	return mangos.ProtoPush
    76  }
    77  
    78  func (*push) PeerNumber() uint16 {
    79  	return mangos.ProtoPull
    80  }
    81  
    82  func (*push) Name() string {
    83  	return "push"
    84  }
    85  
    86  func (*push) PeerName() string {
    87  	return "pull"
    88  }
    89  
    90  func (x *push) AddEndpoint(ep mangos.Endpoint) {
    91  	pe := &pushEp{ep: ep, cq: make(chan struct{})}
    92  	x.Lock()
    93  	x.eps[ep.GetID()] = pe
    94  	x.Unlock()
    95  	x.w.Add()
    96  	go x.sender(pe)
    97  	go mangos.NullRecv(ep)
    98  }
    99  
   100  func (x *push) RemoveEndpoint(ep mangos.Endpoint) {
   101  	id := ep.GetID()
   102  	x.Lock()
   103  	pe := x.eps[id]
   104  	delete(x.eps, id)
   105  	x.Unlock()
   106  	if pe != nil {
   107  		close(pe.cq)
   108  	}
   109  }
   110  
   111  func (x *push) SetOption(name string, v interface{}) error {
   112  	var ok bool
   113  	switch name {
   114  	case mangos.OptionRaw:
   115  		if x.raw, ok = v.(bool); !ok {
   116  			return mangos.ErrBadValue
   117  		}
   118  		return nil
   119  	default:
   120  		return mangos.ErrBadOption
   121  	}
   122  }
   123  
   124  func (x *push) GetOption(name string) (interface{}, error) {
   125  	switch name {
   126  	case mangos.OptionRaw:
   127  		return x.raw, nil
   128  	default:
   129  		return nil, mangos.ErrBadOption
   130  	}
   131  }
   132  
   133  // NewSocket allocates a new Socket using the PUSH protocol.
   134  func NewSocket() (mangos.Socket, error) {
   135  	return mangos.MakeSocket(&push{}), nil
   136  }