github.com/gdamore/mangos@v1.4.0/protocol/sub/sub.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 sub implements the SUB protocol.  This protocol receives messages
    16  // from publishers (PUB peers).  The messages are filtered based on
    17  // subscription, such that only subscribed messages (see OptionSubscribe) are
    18  // received.
    19  //
    20  // Note that in order to receive any messages, at least one subscription must
    21  // be present.  If no subscription is present (the default state), receive
    22  // operations will block forever.
    23  package sub
    24  
    25  import (
    26  	"bytes"
    27  	"sync"
    28  	"time"
    29  
    30  	"nanomsg.org/go-mangos"
    31  )
    32  
    33  type sub struct {
    34  	sock mangos.ProtocolSocket
    35  	subs [][]byte
    36  	raw  bool
    37  	sync.Mutex
    38  }
    39  
    40  func (s *sub) Init(sock mangos.ProtocolSocket) {
    41  	s.sock = sock
    42  	s.subs = [][]byte{}
    43  	s.sock.SetSendError(mangos.ErrProtoOp)
    44  }
    45  
    46  func (*sub) Shutdown(time.Time) {} // No sender to drain.
    47  
    48  func (s *sub) receiver(ep mangos.Endpoint) {
    49  
    50  	rq := s.sock.RecvChannel()
    51  	cq := s.sock.CloseChannel()
    52  
    53  	for {
    54  		var matched = false
    55  
    56  		m := ep.RecvMsg()
    57  		if m == nil {
    58  			return
    59  		}
    60  
    61  		s.Lock()
    62  		for _, sub := range s.subs {
    63  			if bytes.HasPrefix(m.Body, sub) {
    64  				// Matched, send it up.  Best effort.
    65  				matched = true
    66  				break
    67  			}
    68  		}
    69  		s.Unlock()
    70  
    71  		if !matched {
    72  			m.Free()
    73  			continue
    74  		}
    75  
    76  		select {
    77  		case rq <- m:
    78  		case <-cq:
    79  			m.Free()
    80  			return
    81  		default: // no room, drop it
    82  			m.Free()
    83  		}
    84  	}
    85  }
    86  
    87  func (*sub) Number() uint16 {
    88  	return mangos.ProtoSub
    89  }
    90  
    91  func (*sub) PeerNumber() uint16 {
    92  	return mangos.ProtoPub
    93  }
    94  
    95  func (*sub) Name() string {
    96  	return "sub"
    97  }
    98  
    99  func (*sub) PeerName() string {
   100  	return "pub"
   101  }
   102  
   103  func (s *sub) AddEndpoint(ep mangos.Endpoint) {
   104  	go s.receiver(ep)
   105  }
   106  
   107  func (*sub) RemoveEndpoint(mangos.Endpoint) {}
   108  
   109  func (s *sub) SetOption(name string, value interface{}) error {
   110  	s.Lock()
   111  	defer s.Unlock()
   112  
   113  	var vb []byte
   114  	var ok bool
   115  
   116  	// Check names first, because type check below is only valid for
   117  	// subscription options.
   118  	switch name {
   119  	case mangos.OptionRaw:
   120  		if s.raw, ok = value.(bool); !ok {
   121  			return mangos.ErrBadValue
   122  		}
   123  		return nil
   124  	case mangos.OptionSubscribe:
   125  	case mangos.OptionUnsubscribe:
   126  	default:
   127  		return mangos.ErrBadOption
   128  	}
   129  
   130  	switch v := value.(type) {
   131  	case []byte:
   132  		vb = v
   133  	case string:
   134  		vb = []byte(v)
   135  	default:
   136  		return mangos.ErrBadValue
   137  	}
   138  	switch name {
   139  	case mangos.OptionSubscribe:
   140  		for _, sub := range s.subs {
   141  			if bytes.Equal(sub, vb) {
   142  				// Already present
   143  				return nil
   144  			}
   145  		}
   146  		s.subs = append(s.subs, vb)
   147  		return nil
   148  
   149  	case mangos.OptionUnsubscribe:
   150  		for i, sub := range s.subs {
   151  			if bytes.Equal(sub, vb) {
   152  				s.subs[i] = s.subs[len(s.subs)-1]
   153  				s.subs = s.subs[:len(s.subs)-1]
   154  				return nil
   155  			}
   156  		}
   157  		// Subscription not present
   158  		return mangos.ErrBadValue
   159  
   160  	default:
   161  		return mangos.ErrBadOption
   162  	}
   163  }
   164  
   165  func (s *sub) GetOption(name string) (interface{}, error) {
   166  	switch name {
   167  	case mangos.OptionRaw:
   168  		return s.raw, nil
   169  	default:
   170  		return nil, mangos.ErrBadOption
   171  	}
   172  }
   173  
   174  // NewSocket allocates a new Socket using the SUB protocol.
   175  func NewSocket() (mangos.Socket, error) {
   176  	return mangos.MakeSocket(&sub{}), nil
   177  }