github.com/grahambrereton-form3/tilt@v0.10.18/internal/hud/server/websocket.go (about)

     1  package server
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"io"
     7  	"net/http"
     8  	"sync/atomic"
     9  
    10  	"github.com/grpc-ecosystem/grpc-gateway/runtime"
    11  
    12  	"github.com/windmilleng/tilt/internal/hud/webview"
    13  	"github.com/windmilleng/tilt/internal/store"
    14  	"github.com/windmilleng/tilt/pkg/logger"
    15  
    16  	"github.com/gorilla/websocket"
    17  )
    18  
    19  var upgrader = websocket.Upgrader{
    20  	ReadBufferSize:  1024,
    21  	WriteBufferSize: 1024,
    22  }
    23  
    24  type WebsocketSubscriber struct {
    25  	conn       WebsocketConn
    26  	streamDone chan bool
    27  }
    28  
    29  type WebsocketConn interface {
    30  	NextReader() (int, io.Reader, error)
    31  	Close() error
    32  	WriteJSON(v interface{}) error
    33  	NextWriter(messageType int) (io.WriteCloser, error)
    34  }
    35  
    36  var _ WebsocketConn = &websocket.Conn{}
    37  
    38  func NewWebsocketSubscriber(conn WebsocketConn) WebsocketSubscriber {
    39  	return WebsocketSubscriber{
    40  		conn:       conn,
    41  		streamDone: make(chan bool, 0),
    42  	}
    43  }
    44  
    45  func (ws WebsocketSubscriber) TearDown(ctx context.Context) {
    46  	_ = ws.conn.Close()
    47  }
    48  
    49  // Should be called exactly once. Consumes messages until the socket closes.
    50  func (ws WebsocketSubscriber) Stream(ctx context.Context, store *store.Store) {
    51  	go func() {
    52  		// No-op consumption of all control messages, as recommended here:
    53  		// https://godoc.org/github.com/gorilla/websocket#hdr-Control_Messages
    54  		conn := ws.conn
    55  		for {
    56  			if _, _, err := conn.NextReader(); err != nil {
    57  				close(ws.streamDone)
    58  				break
    59  			}
    60  		}
    61  	}()
    62  
    63  	<-ws.streamDone
    64  
    65  	// When we remove ourselves as a subscriber, the Store waits for any outstanding OnChange
    66  	// events to complete, then calls TearDown.
    67  	_ = store.RemoveSubscriber(context.Background(), ws)
    68  }
    69  
    70  func (ws WebsocketSubscriber) OnChange(ctx context.Context, s store.RStore) {
    71  	state := s.RLockState()
    72  	view, err := webview.StateToProtoView(state)
    73  	if err != nil {
    74  		logger.Get(ctx).Infof("error converting view to proto for websocket: %v", err)
    75  		return
    76  	}
    77  
    78  	if view.NeedsAnalyticsNudge && !state.AnalyticsNudgeSurfaced {
    79  		// If we're showing the nudge and no one's told the engine
    80  		// state about it yet... tell the engine state.
    81  		s.Dispatch(store.AnalyticsNudgeSurfacedAction{})
    82  	}
    83  	s.RUnlockState()
    84  
    85  	jsEncoder := &runtime.JSONPb{OrigName: false, EmitDefaults: true}
    86  	w, err := ws.conn.NextWriter(websocket.TextMessage)
    87  	if err != nil {
    88  		logger.Get(ctx).Verbosef("getting writer: %v", err)
    89  	}
    90  	defer func() {
    91  		err := w.Close()
    92  		if err != nil {
    93  			logger.Get(ctx).Verbosef("error closing websocket: %v", err)
    94  		}
    95  	}()
    96  
    97  	err = jsEncoder.NewEncoder(w).Encode(view)
    98  	if err != nil {
    99  		logger.Get(ctx).Verbosef("sending webview data: %v", err)
   100  	}
   101  }
   102  
   103  func (s *HeadsUpServer) ViewWebsocket(w http.ResponseWriter, req *http.Request) {
   104  	conn, err := upgrader.Upgrade(w, req, nil)
   105  	if err != nil {
   106  		http.Error(w, fmt.Sprintf("Error upgrading websocket: %v", err), http.StatusInternalServerError)
   107  		return
   108  	}
   109  
   110  	atomic.AddInt32(&s.numWebsocketConns, 1)
   111  	ws := NewWebsocketSubscriber(conn)
   112  
   113  	// TODO(nick): Handle clean shutdown when the server shuts down
   114  	ctx := context.TODO()
   115  
   116  	// Fire a fake OnChange event to initialize the connection.
   117  	ws.OnChange(ctx, s.store)
   118  	s.store.AddSubscriber(ctx, ws)
   119  
   120  	ws.Stream(ctx, s.store)
   121  	atomic.AddInt32(&s.numWebsocketConns, -1)
   122  }
   123  
   124  var _ store.TearDowner = WebsocketSubscriber{}