func (r *Request) FromHTTP(req *http.Request) error {
if req == nil {
return errors.New("natsproxy: Request cannot be nil")
}
isWebSock := IsWebSocketRequest(req)
wsID := ""
if isWebSock {
wsID = nuid.Next()
}
buf := bytes.NewBuffer(r.Body)
buf.Reset()
if req.Body != nil {
if _, err := io.Copy(buf, req.Body); err != nil {
return err
}
if err := req.Body.Close(); err != nil {
return err
}
}
headerMap := copyMap(map[string][]string(req.Header))
r.URL = req.URL.String()
r.Method = req.Method
r.Header = headerMap
r.RemoteAddr = req.RemoteAddr
r.WebSocketID = wsID
r.Body = buf.Bytes()
return nil
}
// NewBenchmark initializes a Benchmark. After creating a bench call AddSubSample/AddPubSample.
// When done collecting samples, call EndBenchmark
func NewBenchmark(name string, subCnt, pubCnt int) *Benchmark {
bm := Benchmark{Name: name, RunID: nuid.Next()}
bm.Subs = NewSampleGroup()
bm.Pubs = NewSampleGroup()
bm.subChannel = make(chan *Sample, subCnt)
bm.pubChannel = make(chan *Sample, pubCnt)
return &bm
}
func (sc *conn) publishAsync(subject string, data []byte, ah AckHandler, ch chan error) (string, error) {
a := &ack{ah: ah, ch: ch}
sc.Lock()
if sc.nc == nil {
sc.Unlock()
return "", ErrConnectionClosed
}
subj := sc.pubPrefix + "." + subject
// This is only what we need from PubMsg in the timer below,
// so do this so that pe doesn't escape (and we same on new object)
peGUID := nuid.Next()
pe := &pb.PubMsg{ClientID: sc.clientID, Guid: peGUID, Subject: subject, Data: data}
b, _ := pe.Marshal()
// Map ack to guid.
sc.pubAckMap[peGUID] = a
// snapshot
ackSubject := sc.ackSubject
ackTimeout := sc.opts.AckTimeout
pac := sc.pubAckChan
sc.Unlock()
// Use the buffered channel to control the number of outstanding acks.
pac <- struct{}{}
err := sc.nc.PublishRequest(subj, ackSubject, b)
if err != nil {
sc.removeAck(peGUID)
return "", err
}
// Setup the timer for expiration.
sc.Lock()
a.t = time.AfterFunc(ackTimeout, func() {
sc.removeAck(peGUID)
if a.ah != nil {
ah(peGUID, ErrTimeout)
} else if a.ch != nil {
a.ch <- ErrTimeout
}
})
sc.Unlock()
return peGUID, nil
}
// Use nuid.
func genID() string {
return nuid.Next()
}
// Connect will form a connection to the NATS Streaming subsystem.
func Connect(stanClusterID, clientID string, options ...Option) (Conn, error) {
// Process Options
c := conn{clientID: clientID, opts: DefaultOptions}
for _, opt := range options {
if err := opt(&c.opts); err != nil {
return nil, err
}
}
// Check if the user has provided a connection as an option
c.nc = c.opts.NatsConn
// Create a NATS connection if it doesn't exist.
if c.nc == nil {
nc, err := nats.Connect(c.opts.NatsURL)
if err != nil {
return nil, err
}
c.nc = nc
c.ncOwned = true
} else if !c.nc.IsConnected() {
// Bail if the custom NATS connection is disconnected
return nil, ErrBadConnection
}
// Create a heartbeat inbox
hbInbox := nats.NewInbox()
var err error
if c.hbSubscription, err = c.nc.Subscribe(hbInbox, c.processHeartBeat); err != nil {
c.Close()
return nil, err
}
// Send Request to discover the cluster
discoverSubject := c.opts.DiscoverPrefix + "." + stanClusterID
req := &pb.ConnectRequest{ClientID: clientID, HeartbeatInbox: hbInbox}
b, _ := req.Marshal()
reply, err := c.nc.Request(discoverSubject, b, c.opts.ConnectTimeout)
if err != nil {
c.Close()
if err == nats.ErrTimeout {
return nil, ErrConnectReqTimeout
}
return nil, err
}
// Process the response, grab server pubPrefix
cr := &pb.ConnectResponse{}
err = cr.Unmarshal(reply.Data)
if err != nil {
c.Close()
return nil, err
}
if cr.Error != "" {
c.Close()
return nil, errors.New(cr.Error)
}
// Capture cluster configuration endpoints to publish and subscribe/unsubscribe.
c.pubPrefix = cr.PubPrefix
c.subRequests = cr.SubRequests
c.unsubRequests = cr.UnsubRequests
c.closeRequests = cr.CloseRequests
// Setup the ACK subscription
c.ackSubject = DefaultACKPrefix + "." + nuid.Next()
if c.ackSubscription, err = c.nc.Subscribe(c.ackSubject, c.processAck); err != nil {
c.Close()
return nil, err
}
c.ackSubscription.SetPendingLimits(1024*1024, 32*1024*1024)
c.pubAckMap = make(map[string]*ack)
// Create Subscription map
c.subMap = make(map[string]*subscription)
c.pubAckChan = make(chan struct{}, c.opts.MaxPubAcksInflight)
// Attach a finalizer
runtime.SetFinalizer(&c, func(sc *conn) { sc.Close() })
return &c, nil
}