-
Notifications
You must be signed in to change notification settings - Fork 6
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Refactor job controller and websocket logic. (#150)
- Loading branch information
1 parent
4370d14
commit 5c1acd9
Showing
27 changed files
with
556 additions
and
1,048 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file was deleted.
Oops, something went wrong.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,63 @@ | ||
| package websocket | ||
|
|
||
| import ( | ||
| "net/http" | ||
|
|
||
| "github.com/gorilla/websocket" | ||
| "github.com/pkg/errors" | ||
| "github.com/runatlantis/atlantis/server/logging" | ||
| ) | ||
|
|
||
| // PartitionKeyGenerator generates partition keys for the multiplexor | ||
| type PartitionKeyGenerator interface { | ||
| Generate(r *http.Request) (string, error) | ||
| } | ||
|
|
||
| // PartitionRegistry is the registry holding each partition | ||
| // and is responsible for registering/deregistering new buffers | ||
| type PartitionRegistry interface { | ||
| Register(key string, buffer chan string) | ||
| Deregister(key string, buffer chan string) | ||
| } | ||
|
|
||
| // Multiplexor is responsible for handling the data transfer between the storage layer | ||
| // and the registry. Note this is still a WIP as right now the registry is assumed to handle | ||
| // everything. | ||
| type Multiplexor struct { | ||
| writer *Writer | ||
| keyGenerator PartitionKeyGenerator | ||
| registry PartitionRegistry | ||
| } | ||
|
|
||
| func NewMultiplexor(log logging.SimpleLogging, keyGenerator PartitionKeyGenerator, registry PartitionRegistry) *Multiplexor { | ||
| upgrader := websocket.Upgrader{} | ||
| upgrader.CheckOrigin = func(r *http.Request) bool { return true } | ||
| return &Multiplexor{ | ||
| writer: &Writer{ | ||
| upgrader: upgrader, | ||
| log: log, | ||
| }, | ||
| keyGenerator: keyGenerator, | ||
| registry: registry, | ||
| } | ||
| } | ||
|
|
||
| // Handle should be called for a given websocket request. It blocks | ||
| // while writing to the websocket until the buffer is closed. | ||
| func (m *Multiplexor) Handle(w http.ResponseWriter, r *http.Request) error { | ||
| key, err := m.keyGenerator.Generate(r) | ||
|
|
||
| if err != nil { | ||
| return errors.Wrapf(err, "generating partition key") | ||
| } | ||
|
|
||
| // Buffer size set to 1000 to ensure messages get queued. | ||
| // TODO: make buffer size configurable | ||
| buffer := make(chan string, 1000) | ||
|
|
||
| // spinning up a goroutine for this since we are attempting to block on the read side. | ||
| go m.registry.Register(key, buffer) | ||
| defer m.registry.Deregister(key, buffer) | ||
|
|
||
| return errors.Wrapf(m.writer.Write(w, r, buffer), "writing to ws %s", key) | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,70 @@ | ||
| package websocket | ||
|
|
||
| import ( | ||
| "net/http" | ||
|
|
||
| "github.com/gorilla/websocket" | ||
| "github.com/pkg/errors" | ||
| "github.com/runatlantis/atlantis/server/logging" | ||
| ) | ||
|
|
||
| func NewWriter(log logging.SimpleLogging) *Writer { | ||
| upgrader := websocket.Upgrader{} | ||
| upgrader.CheckOrigin = func(r *http.Request) bool { return true } | ||
| return &Writer{ | ||
| upgrader: upgrader, | ||
| log: log, | ||
| } | ||
| } | ||
|
|
||
| type Writer struct { | ||
| upgrader websocket.Upgrader | ||
|
|
||
| //TODO: Remove dependency on atlantis logger here if we upstream this. | ||
| log logging.SimpleLogging | ||
| } | ||
|
|
||
| func (w *Writer) Write(rw http.ResponseWriter, r *http.Request, input chan string) error { | ||
| conn, err := w.upgrader.Upgrade(rw, r, nil) | ||
|
|
||
| if err != nil { | ||
| return errors.Wrap(err, "upgrading websocket connection") | ||
| } | ||
|
|
||
| conn.SetCloseHandler(func(code int, text string) error { | ||
| // Close the channnel after websocket connection closed. | ||
| // Will gracefully exit the ProjectCommandOutputHandler.Register() call and cleanup. | ||
| // is it good practice to close at the receiver? Probably not, we should figure out a better | ||
| // way to handle this case | ||
| close(input) | ||
| return nil | ||
| }) | ||
|
|
||
| // Add a reader goroutine to listen for socket.close() events. | ||
| go w.setReadHandler(conn) | ||
|
|
||
| // block on reading our input channel | ||
| for msg := range input { | ||
| if err := conn.WriteMessage(websocket.BinaryMessage, []byte("\r"+msg+"\n")); err != nil { | ||
| w.log.Warn("Failed to write ws message: %s", err) | ||
| return err | ||
| } | ||
| } | ||
|
|
||
| return nil | ||
| } | ||
|
|
||
| func (w *Writer) setReadHandler(c *websocket.Conn) { | ||
| for { | ||
| _, _, err := c.ReadMessage() | ||
| if err != nil { | ||
| // CloseGoingAway (1001) when a browser tab is closed. | ||
| // Expected behaviour since we have a CloseHandler(), log warning if not a CloseGoingAway | ||
| if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway) { | ||
| w.log.Warn("Failed to read WS message: %s", err) | ||
| } | ||
| return | ||
| } | ||
| } | ||
|
|
||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.