github.com/pf-qiu/concourse/v6@v6.7.3-0.20201207032516-1f455d73275f/atc/api/buildserver/eventhandler.go (about)

     1  package buildserver
     2  
     3  import (
     4  	"encoding/json"
     5  	"fmt"
     6  	"io"
     7  	"net/http"
     8  
     9  	"code.cloudfoundry.org/lager"
    10  	"github.com/pf-qiu/concourse/v6/atc/db"
    11  	"github.com/vito/go-sse/sse"
    12  )
    13  
    14  const ProtocolVersionHeader = "X-ATC-Stream-Version"
    15  const CurrentProtocolVersion = "2.0"
    16  
    17  func NewEventHandler(logger lager.Logger, build db.Build) http.Handler {
    18  	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    19  		var eventID uint = 0
    20  		if r.Header.Get("Last-Event-ID") != "" {
    21  			startString := r.Header.Get("Last-Event-ID")
    22  			_, err := fmt.Sscanf(startString, "%d", &eventID)
    23  			if err != nil {
    24  				logger.Info("failed-to-parse-last-event-id", lager.Data{"last-event-id": startString})
    25  				w.WriteHeader(http.StatusBadRequest)
    26  				return
    27  			}
    28  
    29  			eventID++
    30  		}
    31  
    32  		w.Header().Add("Content-Type", "text/event-stream; charset=utf-8")
    33  		w.Header().Add("Cache-Control", "no-cache, no-store, must-revalidate")
    34  		w.Header().Add("X-Accel-Buffering", "no")
    35  		w.Header().Add(ProtocolVersionHeader, CurrentProtocolVersion)
    36  
    37  		writer := eventWriter{
    38  			responseWriter:  w,
    39  			responseFlusher: w.(http.Flusher),
    40  		}
    41  
    42  		events, err := build.Events(eventID)
    43  		if err != nil {
    44  			logger.Error("failed-to-get-build-events", err, lager.Data{"build-id": build.ID(), "start": eventID})
    45  			w.WriteHeader(http.StatusInternalServerError)
    46  			return
    47  		}
    48  
    49  		defer db.Close(events)
    50  
    51  		for {
    52  			logger = logger.WithData(lager.Data{"id": eventID})
    53  
    54  			ev, err := events.Next()
    55  			if err != nil {
    56  				if err == db.ErrEndOfBuildEventStream {
    57  					err := writer.WriteEnd(eventID)
    58  					if err != nil {
    59  						logger.Info("failed-to-write-end", lager.Data{"error": err.Error()})
    60  						return
    61  					}
    62  
    63  					<-r.Context().Done()
    64  				} else {
    65  					logger.Error("failed-to-get-next-build-event", err)
    66  					return
    67  				}
    68  
    69  				return
    70  			}
    71  
    72  			err = writer.WriteEvent(eventID, ev)
    73  			if err != nil {
    74  				logger.Info("failed-to-write-event", lager.Data{"error": err.Error()})
    75  				return
    76  			}
    77  
    78  			eventID++
    79  		}
    80  	})
    81  }
    82  
    83  type eventWriter struct {
    84  	responseWriter  io.Writer
    85  	responseFlusher http.Flusher
    86  }
    87  
    88  func (writer eventWriter) WriteEvent(id uint, envelope interface{}) error {
    89  	payload, err := json.Marshal(envelope)
    90  	if err != nil {
    91  		return err
    92  	}
    93  
    94  	err = sse.Event{
    95  		ID:   fmt.Sprintf("%d", id),
    96  		Name: "event",
    97  		Data: payload,
    98  	}.Write(writer.responseWriter)
    99  	if err != nil {
   100  		return err
   101  	}
   102  
   103  	writer.responseFlusher.Flush()
   104  
   105  	return nil
   106  }
   107  
   108  func (writer eventWriter) WriteEnd(id uint) error {
   109  	err := sse.Event{
   110  		ID:   fmt.Sprintf("%d", id),
   111  		Name: "end",
   112  	}.Write(writer.responseWriter)
   113  	if err != nil {
   114  		return err
   115  	}
   116  
   117  	writer.responseFlusher.Flush()
   118  
   119  	return nil
   120  }