Initial QSfera import
This commit is contained in:
Курнат Андрей
+390
@@ -0,0 +1,390 @@
|
||||
/* This Source Code Form is subject to the terms of the Mozilla Public
|
||||
* License, v. 2.0. If a copy of the MPL was not distributed with this
|
||||
* file, You can obtain one at http://mozilla.org/MPL/2.0/. */
|
||||
|
||||
package sse
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"gopkg.in/cenkalti/backoff.v1"
|
||||
)
|
||||
|
||||
var (
|
||||
headerID = []byte("id:")
|
||||
headerData = []byte("data:")
|
||||
headerEvent = []byte("event:")
|
||||
headerRetry = []byte("retry:")
|
||||
)
|
||||
|
||||
func ClientMaxBufferSize(s int) func(c *Client) {
|
||||
return func(c *Client) {
|
||||
c.maxBufferSize = s
|
||||
}
|
||||
}
|
||||
|
||||
// ConnCallback defines a function to be called on a particular connection event
|
||||
type ConnCallback func(c *Client)
|
||||
|
||||
// ResponseValidator validates a response
|
||||
type ResponseValidator func(c *Client, resp *http.Response) error
|
||||
|
||||
// Client handles an incoming server stream
|
||||
type Client struct {
|
||||
Retry time.Time
|
||||
ReconnectStrategy backoff.BackOff
|
||||
disconnectcb ConnCallback
|
||||
connectedcb ConnCallback
|
||||
subscribed map[chan *Event]chan struct{}
|
||||
Headers map[string]string
|
||||
ReconnectNotify backoff.Notify
|
||||
ResponseValidator ResponseValidator
|
||||
Connection *http.Client
|
||||
URL string
|
||||
LastEventID atomic.Value // []byte
|
||||
maxBufferSize int
|
||||
mu sync.Mutex
|
||||
EncodingBase64 bool
|
||||
Connected bool
|
||||
}
|
||||
|
||||
// NewClient creates a new client
|
||||
func NewClient(url string, opts ...func(c *Client)) *Client {
|
||||
c := &Client{
|
||||
URL: url,
|
||||
Connection: &http.Client{},
|
||||
Headers: make(map[string]string),
|
||||
subscribed: make(map[chan *Event]chan struct{}),
|
||||
maxBufferSize: 1 << 16,
|
||||
}
|
||||
|
||||
for _, opt := range opts {
|
||||
opt(c)
|
||||
}
|
||||
|
||||
return c
|
||||
}
|
||||
|
||||
// Subscribe to a data stream
|
||||
func (c *Client) Subscribe(stream string, handler func(msg *Event)) error {
|
||||
return c.SubscribeWithContext(context.Background(), stream, handler)
|
||||
}
|
||||
|
||||
// SubscribeWithContext to a data stream with context
|
||||
func (c *Client) SubscribeWithContext(ctx context.Context, stream string, handler func(msg *Event)) error {
|
||||
operation := func() error {
|
||||
resp, err := c.request(ctx, stream)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if validator := c.ResponseValidator; validator != nil {
|
||||
err = validator(c, resp)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
} else if resp.StatusCode != 200 {
|
||||
resp.Body.Close()
|
||||
return fmt.Errorf("could not connect to stream: %s", http.StatusText(resp.StatusCode))
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
reader := NewEventStreamReader(resp.Body, c.maxBufferSize)
|
||||
eventChan, errorChan := c.startReadLoop(reader)
|
||||
|
||||
for {
|
||||
select {
|
||||
case err = <-errorChan:
|
||||
return err
|
||||
case msg := <-eventChan:
|
||||
handler(msg)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Apply user specified reconnection strategy or default to standard NewExponentialBackOff() reconnection method
|
||||
var err error
|
||||
if c.ReconnectStrategy != nil {
|
||||
err = backoff.RetryNotify(operation, c.ReconnectStrategy, c.ReconnectNotify)
|
||||
} else {
|
||||
err = backoff.RetryNotify(operation, backoff.NewExponentialBackOff(), c.ReconnectNotify)
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
// SubscribeChan sends all events to the provided channel
|
||||
func (c *Client) SubscribeChan(stream string, ch chan *Event) error {
|
||||
return c.SubscribeChanWithContext(context.Background(), stream, ch)
|
||||
}
|
||||
|
||||
// SubscribeChanWithContext sends all events to the provided channel with context
|
||||
func (c *Client) SubscribeChanWithContext(ctx context.Context, stream string, ch chan *Event) error {
|
||||
var connected bool
|
||||
errch := make(chan error)
|
||||
c.mu.Lock()
|
||||
c.subscribed[ch] = make(chan struct{})
|
||||
c.mu.Unlock()
|
||||
|
||||
operation := func() error {
|
||||
resp, err := c.request(ctx, stream)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if validator := c.ResponseValidator; validator != nil {
|
||||
err = validator(c, resp)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
} else if resp.StatusCode != 200 {
|
||||
resp.Body.Close()
|
||||
return fmt.Errorf("could not connect to stream: %s", http.StatusText(resp.StatusCode))
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if !connected {
|
||||
// Notify connect
|
||||
errch <- nil
|
||||
connected = true
|
||||
}
|
||||
|
||||
reader := NewEventStreamReader(resp.Body, c.maxBufferSize)
|
||||
eventChan, errorChan := c.startReadLoop(reader)
|
||||
|
||||
for {
|
||||
var msg *Event
|
||||
// Wait for message to arrive or exit
|
||||
select {
|
||||
case <-c.subscribed[ch]:
|
||||
return nil
|
||||
case err = <-errorChan:
|
||||
return err
|
||||
case msg = <-eventChan:
|
||||
}
|
||||
|
||||
// Wait for message to be sent or exit
|
||||
if msg != nil {
|
||||
select {
|
||||
case <-c.subscribed[ch]:
|
||||
return nil
|
||||
case ch <- msg:
|
||||
// message sent
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
go func() {
|
||||
defer c.cleanup(ch)
|
||||
// Apply user specified reconnection strategy or default to standard NewExponentialBackOff() reconnection method
|
||||
var err error
|
||||
if c.ReconnectStrategy != nil {
|
||||
err = backoff.RetryNotify(operation, c.ReconnectStrategy, c.ReconnectNotify)
|
||||
} else {
|
||||
err = backoff.RetryNotify(operation, backoff.NewExponentialBackOff(), c.ReconnectNotify)
|
||||
}
|
||||
|
||||
// channel closed once connected
|
||||
if err != nil && !connected {
|
||||
errch <- err
|
||||
}
|
||||
}()
|
||||
err := <-errch
|
||||
close(errch)
|
||||
return err
|
||||
}
|
||||
|
||||
func (c *Client) startReadLoop(reader *EventStreamReader) (chan *Event, chan error) {
|
||||
outCh := make(chan *Event)
|
||||
erChan := make(chan error)
|
||||
go c.readLoop(reader, outCh, erChan)
|
||||
return outCh, erChan
|
||||
}
|
||||
|
||||
func (c *Client) readLoop(reader *EventStreamReader, outCh chan *Event, erChan chan error) {
|
||||
for {
|
||||
// Read each new line and process the type of event
|
||||
event, err := reader.ReadEvent()
|
||||
if err != nil {
|
||||
if err == io.EOF {
|
||||
erChan <- nil
|
||||
return
|
||||
}
|
||||
// run user specified disconnect function
|
||||
if c.disconnectcb != nil {
|
||||
c.Connected = false
|
||||
c.disconnectcb(c)
|
||||
}
|
||||
erChan <- err
|
||||
return
|
||||
}
|
||||
|
||||
if !c.Connected && c.connectedcb != nil {
|
||||
c.Connected = true
|
||||
c.connectedcb(c)
|
||||
}
|
||||
|
||||
// If we get an error, ignore it.
|
||||
var msg *Event
|
||||
if msg, err = c.processEvent(event); err == nil {
|
||||
if len(msg.ID) > 0 {
|
||||
c.LastEventID.Store(msg.ID)
|
||||
} else {
|
||||
msg.ID, _ = c.LastEventID.Load().([]byte)
|
||||
}
|
||||
|
||||
// Send downstream if the event has something useful
|
||||
if msg.hasContent() {
|
||||
outCh <- msg
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// SubscribeRaw to an sse endpoint
|
||||
func (c *Client) SubscribeRaw(handler func(msg *Event)) error {
|
||||
return c.Subscribe("", handler)
|
||||
}
|
||||
|
||||
// SubscribeRawWithContext to an sse endpoint with context
|
||||
func (c *Client) SubscribeRawWithContext(ctx context.Context, handler func(msg *Event)) error {
|
||||
return c.SubscribeWithContext(ctx, "", handler)
|
||||
}
|
||||
|
||||
// SubscribeChanRaw sends all events to the provided channel
|
||||
func (c *Client) SubscribeChanRaw(ch chan *Event) error {
|
||||
return c.SubscribeChan("", ch)
|
||||
}
|
||||
|
||||
// SubscribeChanRawWithContext sends all events to the provided channel with context
|
||||
func (c *Client) SubscribeChanRawWithContext(ctx context.Context, ch chan *Event) error {
|
||||
return c.SubscribeChanWithContext(ctx, "", ch)
|
||||
}
|
||||
|
||||
// Unsubscribe unsubscribes a channel
|
||||
func (c *Client) Unsubscribe(ch chan *Event) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
|
||||
if c.subscribed[ch] != nil {
|
||||
c.subscribed[ch] <- struct{}{}
|
||||
}
|
||||
}
|
||||
|
||||
// OnDisconnect specifies the function to run when the connection disconnects
|
||||
func (c *Client) OnDisconnect(fn ConnCallback) {
|
||||
c.disconnectcb = fn
|
||||
}
|
||||
|
||||
// OnConnect specifies the function to run when the connection is successful
|
||||
func (c *Client) OnConnect(fn ConnCallback) {
|
||||
c.connectedcb = fn
|
||||
}
|
||||
|
||||
func (c *Client) request(ctx context.Context, stream string) (*http.Response, error) {
|
||||
req, err := http.NewRequest("GET", c.URL, nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req = req.WithContext(ctx)
|
||||
|
||||
// Setup request, specify stream to connect to
|
||||
if stream != "" {
|
||||
query := req.URL.Query()
|
||||
query.Add("stream", stream)
|
||||
req.URL.RawQuery = query.Encode()
|
||||
}
|
||||
|
||||
req.Header.Set("Cache-Control", "no-cache")
|
||||
req.Header.Set("Accept", "text/event-stream")
|
||||
req.Header.Set("Connection", "keep-alive")
|
||||
|
||||
lastID, exists := c.LastEventID.Load().([]byte)
|
||||
if exists && lastID != nil {
|
||||
req.Header.Set("Last-Event-ID", string(lastID))
|
||||
}
|
||||
|
||||
// Add user specified headers
|
||||
for k, v := range c.Headers {
|
||||
req.Header.Set(k, v)
|
||||
}
|
||||
|
||||
return c.Connection.Do(req)
|
||||
}
|
||||
|
||||
func (c *Client) processEvent(msg []byte) (event *Event, err error) {
|
||||
var e Event
|
||||
|
||||
if len(msg) < 1 {
|
||||
return nil, errors.New("event message was empty")
|
||||
}
|
||||
|
||||
// Normalize the crlf to lf to make it easier to split the lines.
|
||||
// Split the line by "\n" or "\r", per the spec.
|
||||
for _, line := range bytes.FieldsFunc(msg, func(r rune) bool { return r == '\n' || r == '\r' }) {
|
||||
switch {
|
||||
case bytes.HasPrefix(line, headerID):
|
||||
e.ID = append([]byte(nil), trimHeader(len(headerID), line)...)
|
||||
case bytes.HasPrefix(line, headerData):
|
||||
// The spec allows for multiple data fields per event, concatenated them with "\n".
|
||||
e.Data = append(e.Data[:], append(trimHeader(len(headerData), line), byte('\n'))...)
|
||||
// The spec says that a line that simply contains the string "data" should be treated as a data field with an empty body.
|
||||
case bytes.Equal(line, bytes.TrimSuffix(headerData, []byte(":"))):
|
||||
e.Data = append(e.Data, byte('\n'))
|
||||
case bytes.HasPrefix(line, headerEvent):
|
||||
e.Event = append([]byte(nil), trimHeader(len(headerEvent), line)...)
|
||||
case bytes.HasPrefix(line, headerRetry):
|
||||
e.Retry = append([]byte(nil), trimHeader(len(headerRetry), line)...)
|
||||
default:
|
||||
// Ignore any garbage that doesn't match what we're looking for.
|
||||
}
|
||||
}
|
||||
|
||||
// Trim the last "\n" per the spec.
|
||||
e.Data = bytes.TrimSuffix(e.Data, []byte("\n"))
|
||||
|
||||
if c.EncodingBase64 {
|
||||
buf := make([]byte, base64.StdEncoding.DecodedLen(len(e.Data)))
|
||||
|
||||
n, err := base64.StdEncoding.Decode(buf, e.Data)
|
||||
if err != nil {
|
||||
err = fmt.Errorf("failed to decode event message: %s", err)
|
||||
}
|
||||
e.Data = buf[:n]
|
||||
}
|
||||
return &e, err
|
||||
}
|
||||
|
||||
func (c *Client) cleanup(ch chan *Event) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
|
||||
if c.subscribed[ch] != nil {
|
||||
close(c.subscribed[ch])
|
||||
delete(c.subscribed, ch)
|
||||
}
|
||||
}
|
||||
|
||||
func trimHeader(size int, data []byte) []byte {
|
||||
if data == nil || len(data) < size {
|
||||
return data
|
||||
}
|
||||
|
||||
data = data[size:]
|
||||
// Remove optional leading whitespace
|
||||
if len(data) > 0 && data[0] == 32 {
|
||||
data = data[1:]
|
||||
}
|
||||
// Remove trailing new line
|
||||
if len(data) > 0 && data[len(data)-1] == 10 {
|
||||
data = data[:len(data)-1]
|
||||
}
|
||||
return data
|
||||
}
|
||||
Reference in New Issue
Block a user