├── .gitignore ├── go.mod ├── examples ├── command │ └── command.go └── stderr-highlight │ └── stderr-hightlight.go ├── go.sum ├── README.md ├── iomux_test.go ├── iomux.go └── LICENSE.md /.gitignore: -------------------------------------------------------------------------------- 1 | .idea 2 | .vscode 3 | .DS_Store 4 | .history 5 | -------------------------------------------------------------------------------- /go.mod: -------------------------------------------------------------------------------- 1 | module github.com/netflix/go-iomux 2 | 3 | go 1.18 4 | 5 | require ( 6 | github.com/stretchr/testify v1.8.1 7 | golang.org/x/sys v0.4.0 8 | ) 9 | 10 | require ( 11 | github.com/davecgh/go-spew v1.1.1 // indirect 12 | github.com/pmezard/go-difflib v1.0.0 // indirect 13 | gopkg.in/yaml.v3 v3.0.1 // indirect 14 | ) 15 | -------------------------------------------------------------------------------- /examples/command/command.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "encoding/binary" 5 | "errors" 6 | "github.com/netflix/go-iomux" 7 | "io" 8 | "os" 9 | "os/exec" 10 | ) 11 | 12 | type OutputType int 13 | 14 | const ( 15 | StdOut OutputType = iota 16 | StdErr 17 | ) 18 | 19 | func main() { 20 | mux := &iomux.Mux[OutputType]{} 21 | defer mux.Close() 22 | cmd := exec.Command("sh", "-c", "echo out1 && echo err1 1>&2 && echo out2") 23 | stdout, err := mux.Tag(StdOut) 24 | if err != nil { 25 | panic(err) 26 | } 27 | stderr, _ := mux.Tag(StdErr) 28 | cmd.Stdout = stdout 29 | cmd.Stderr = stderr 30 | taggedData, err := mux.ReadWhile(func() error { 31 | return cmd.Run() 32 | }) 33 | for _, td := range taggedData { 34 | var w io.Writer 35 | switch td.Tag { 36 | case StdOut: 37 | w = os.Stdout 38 | case StdErr: 39 | w = os.Stderr 40 | } 41 | binary.Write(w, binary.BigEndian, td.Data) 42 | } 43 | if err != nil { 44 | err = errors.Unwrap(err) 45 | if exitError, ok := err.(*exec.ExitError); ok { 46 | os.Exit(exitError.ProcessState.ExitCode()) 47 | } 48 | } 49 | } 50 | -------------------------------------------------------------------------------- /examples/stderr-highlight/stderr-hightlight.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "context" 5 | "encoding/binary" 6 | "github.com/netflix/go-iomux" 7 | "io" 8 | "os" 9 | "os/exec" 10 | ) 11 | 12 | type OutputType int 13 | 14 | const ( 15 | StdOut OutputType = iota 16 | StdErr 17 | ) 18 | 19 | const colorRed = "\033[31m" 20 | const colorReset = "\033[0m" 21 | 22 | func main() { 23 | mux := iomux.NewMuxUnixGram[OutputType]() 24 | defer mux.Close() 25 | cmd := exec.Command("sh", "-c", "echo out1 && echo err1 1>&2 && echo out2") 26 | stdout, err := mux.Tag(StdOut) 27 | if err != nil { 28 | panic(err) 29 | } 30 | stderr, _ := mux.Tag(StdErr) 31 | cmd.Stdout = stdout 32 | cmd.Stderr = stderr 33 | 34 | ctx, cancelFn := context.WithCancel(context.Background()) 35 | cmd.Start() 36 | go func() { 37 | cmd.Wait() 38 | cancelFn() 39 | }() 40 | for { 41 | b, t, err := mux.Read(ctx) 42 | if err != nil { 43 | if err == io.EOF { 44 | break 45 | } 46 | panic(err) 47 | } else { 48 | switch t { 49 | case StdOut: 50 | binary.Write(os.Stdout, binary.BigEndian, b) 51 | case StdErr: 52 | io.WriteString(os.Stderr, colorRed) 53 | binary.Write(os.Stderr, binary.BigEndian, b) 54 | io.WriteString(os.Stderr, colorReset) 55 | } 56 | } 57 | } 58 | } 59 | -------------------------------------------------------------------------------- /go.sum: -------------------------------------------------------------------------------- 1 | github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= 2 | github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= 3 | github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= 4 | github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= 5 | github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= 6 | github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= 7 | github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw= 8 | github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo= 9 | github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= 10 | github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU= 11 | github.com/stretchr/testify v1.8.1 h1:w7B6lhMri9wdJUVmEZPGGhZzrYTPvgJArz7wNPgYKsk= 12 | github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4= 13 | golang.org/x/sys v0.4.0 h1:Zr2JFtRQNX3BCZ8YtxRE9hNJYC8J6I1MVbMg6owUp18= 14 | golang.org/x/sys v0.4.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= 15 | gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= 16 | gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= 17 | gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= 18 | gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= 19 | gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= 20 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # iomux 2 | 3 | [![Go Reference](https://pkg.go.dev/badge/github.com/netflix/go-iomux.svg)](https://pkg.go.dev/github.com/netflix/go-iomux) 4 | 5 | iomux allows multiplexing of file descriptors using any Go supported unix domain network. It makes it possible to distinctly capture `exec.Cmd` stdout/stderr keeping the original output order: 6 | 7 | ``` 8 | mux := &iomux.Mux[OutputType]{} 9 | defer mux.Close() 10 | cmd := exec.Command("sh", "-c", "echo out1 && echo err1 1>&2 && echo out2") 11 | stdout, _ := mux.Tag(StdOut) // errors ignored for brevity 12 | stderr, _ := mux.Tag(StdErr) 13 | cmd.Stdout = stdout 14 | cmd.Stderr = stderr 15 | taggedData, _ := mux.ReadWhile(func() error { 16 | return cmd.Run() 17 | }) 18 | for _, td := range taggedData { 19 | var w io.Writer 20 | switch td.Tag { 21 | case StdOut: 22 | w = os.Stdout 23 | case StdErr: 24 | w = os.Stderr 25 | } 26 | binary.Write(w, binary.BigEndian, td.Data) 27 | } 28 | ``` 29 | 30 | When using `unixgram` networking, ordering is guaranteed (see Limitations below). More can be found in the [examples](examples) directory. 31 | 32 | This module was inspired by Josh Triplett's Rust crate https://github.com/joshtriplett/io-mux/. 33 | 34 | ## Limitations 35 | 36 | On all platforms except macOS the network defaults to `unixgram`. On Linux, `unixgram` behaves like a pipe and will behave exactly as you'd expect, and always see messages in order. On other UNIXes, there is a possibility of different behaviour when [buffers are full](https://docs.rs/io-mux/latest/io_mux/#portability), but it's unlikely a reader will be outpaced. 37 | 38 | Using `unixgram` on macOS when you cannot control the write size of the sender, is impossible due a message size limit of 2048 bytes: 39 | ``` 40 | write /dev/stdout: message too long 41 | ``` 42 | 43 | So on macOs, the default network is `unix`. It is connection oriented, so it doesn't come with the ordering guarantees of `unixgram`. It's possible to see see writes out of order, but on a MacBook Pro M1 0.1ms is the threshold for writes being read out of order, so for real world use cases it's unlikely to be a problem. 44 | 45 | These limitations do not affect the read order of an individual connection, so output for an individual tag is always consistent. If you prefer a different network type, the default can be overridden using the convenience constructors `NewMuxUnix`, `NewMuxUnixGram` and `NewMuxUnixPacket`. 46 | -------------------------------------------------------------------------------- /iomux_test.go: -------------------------------------------------------------------------------- 1 | package iomux 2 | 3 | import ( 4 | "context" 5 | "encoding/binary" 6 | "errors" 7 | "fmt" 8 | "github.com/stretchr/testify/assert" 9 | "golang.org/x/sys/unix" 10 | "io" 11 | "os" 12 | "os/exec" 13 | "testing" 14 | "time" 15 | ) 16 | 17 | var networks = []string{ 18 | "unix", "unixgram", "unixpacket", 19 | } 20 | 21 | const sleepSecs = 0.0001 22 | const sleepDuration = time.Duration(sleepSecs * float64(time.Second)) 23 | 24 | func TestMuxRead(t *testing.T) { 25 | for _, network := range networks { 26 | t.Run(network, func(t *testing.T) { 27 | mux := &Mux[string]{network: network} 28 | t.Cleanup(func() { 29 | mux.Close() 30 | }) 31 | taga, err := mux.Tag("a") 32 | if err != nil { 33 | skipIfProtocolNotSupported(t, err) 34 | assert.Nil(t, err) 35 | } 36 | assert.Nil(t, err) 37 | tagb, _ := mux.Tag("b") 38 | assert.Nil(t, err) 39 | 40 | ctx, cancelFn := context.WithCancel(context.Background()) 41 | io.WriteString(taga, "hello taga") 42 | bytes, tag, err := mux.Read(ctx) 43 | assert.Nil(t, err) 44 | assert.Equal(t, "a", tag) 45 | assert.Equal(t, "hello taga", string(bytes)) 46 | 47 | io.WriteString(tagb, "hello tagb") 48 | bytes, tag, err = mux.Read(ctx) 49 | assert.Nil(t, err) 50 | assert.Equal(t, "b", tag) 51 | assert.Equal(t, "hello tagb", string(bytes)) 52 | 53 | cancelFn() 54 | bytes, tag, err = mux.Read(ctx) 55 | assert.Equal(t, io.EOF, err) 56 | }) 57 | } 58 | } 59 | 60 | func TestMuxReadNoSenders(t *testing.T) { 61 | mux := &Mux[string]{} 62 | t.Cleanup(func() { 63 | mux.Close() 64 | }) 65 | ctx := context.Background() 66 | ctx.Done() 67 | data, tag, err := mux.Read(ctx) 68 | 69 | assert.Nil(t, data) 70 | assert.Equal(t, "", tag) 71 | assert.ErrorIs(t, err, MuxNoConnections) 72 | } 73 | 74 | func TestMuxReadClosed(t *testing.T) { 75 | mux := &Mux[string]{} 76 | mux.Close() 77 | ctx := context.Background() 78 | ctx.Done() 79 | _, _, err := mux.Read(ctx) 80 | 81 | assert.ErrorIs(t, err, MuxClosed) 82 | } 83 | 84 | func TestMuxReadNoData(t *testing.T) { 85 | for _, network := range networks { 86 | t.Run(network, func(t *testing.T) { 87 | mux := &Mux[string]{network: network} 88 | t.Cleanup(func() { 89 | mux.Close() 90 | }) 91 | _, err := mux.Tag("a") 92 | if err != nil { 93 | skipIfProtocolNotSupported(t, err) 94 | assert.Nil(t, err) 95 | } 96 | 97 | ctx, cancelFn := context.WithCancel(context.Background()) 98 | cancelFn() 99 | bytes, tag, err := mux.Read(ctx) 100 | assert.Nil(t, bytes) 101 | assert.Equal(t, "", tag) 102 | assert.Equal(t, io.EOF, err) 103 | }) 104 | } 105 | } 106 | 107 | // More than one active context isn't really an intended use case, but making sure this logic works correctly 108 | func TestMuxMultipleContexts(t *testing.T) { 109 | mux := &Mux[string]{} 110 | t.Cleanup(func() { 111 | mux.Close() 112 | }) 113 | taga, err := mux.Tag("a") 114 | assert.Nil(t, err) 115 | tagb, _ := mux.Tag("b") 116 | 117 | ctx1, cancelFn1 := context.WithCancel(context.Background()) 118 | io.WriteString(taga, "hello taga") 119 | bytes, tag, err := mux.Read(ctx1) 120 | assert.Nil(t, err) 121 | assert.Equal(t, "a", tag) 122 | assert.Equal(t, "hello taga", string(bytes)) 123 | 124 | ctx2, cancelFn2 := context.WithCancel(context.Background()) 125 | io.WriteString(tagb, "hello tagb") 126 | bytes, tag, err = mux.Read(ctx2) 127 | assert.Nil(t, err) 128 | assert.Equal(t, "b", tag) 129 | assert.Equal(t, "hello tagb", string(bytes)) 130 | 131 | cancelFn2() 132 | bytes, tag, err = mux.Read(ctx2) 133 | assert.Equal(t, io.EOF, err) 134 | 135 | cancelFn1() 136 | bytes, tag, err = mux.Read(ctx1) 137 | assert.Equal(t, io.EOF, err) 138 | 139 | assert.Empty(t, mux.recvstate) 140 | } 141 | 142 | func TestMuxReadWhileErr(t *testing.T) { 143 | mux := &Mux[string]{} 144 | t.Cleanup(func() { 145 | mux.Close() 146 | }) 147 | _, err := mux.Tag("a") 148 | assert.Nil(t, err) 149 | 150 | expected := errors.New("this is an error") 151 | _, err = mux.ReadWhile(func() error { 152 | return expected 153 | }) 154 | 155 | assert.ErrorIs(t, expected, err) 156 | } 157 | 158 | func TestMuxTruncatedRead(t *testing.T) { 159 | mux := NewMuxUnix[string]() 160 | t.Cleanup(func() { 161 | mux.Close() 162 | }) 163 | taga, err := mux.Tag("a") 164 | assert.Nil(t, err) 165 | tagb, _ := mux.Tag("b") 166 | 167 | td, err := mux.ReadWhile(func() error { 168 | binary.Write(taga, binary.BigEndian, make([]byte, 256)) 169 | time.Sleep(sleepDuration) 170 | binary.Write(tagb, binary.BigEndian, make([]byte, 10)) 171 | time.Sleep(sleepDuration) 172 | binary.Write(taga, binary.BigEndian, make([]byte, 20)) 173 | return nil 174 | }) 175 | 176 | assert.Equal(t, 3, len(td)) 177 | assert.Equal(t, 256, len(td[0].Data)) 178 | assert.Equal(t, "a", td[0].Tag) 179 | assert.Equal(t, 10, len(td[1].Data)) 180 | assert.Equal(t, "b", td[1].Tag) 181 | assert.Equal(t, 20, len(td[2].Data)) 182 | assert.Equal(t, "a", td[2].Tag) 183 | } 184 | 185 | func skipIfProtocolNotSupported(t *testing.T, err error) { 186 | err = errors.Unwrap(err) 187 | if sys, ok := err.(*os.SyscallError); ok { 188 | if sys.Syscall == "socket" { 189 | err = errors.Unwrap(err) 190 | if err == unix.EPROTONOSUPPORT { 191 | t.Skip("unsupported protocol") 192 | } 193 | } 194 | } 195 | } 196 | 197 | func TestMuxMultiple(t *testing.T) { 198 | for _, network := range networks { 199 | t.Run(network, func(t *testing.T) { 200 | mux := &Mux[string]{network: network} 201 | t.Cleanup(func() { 202 | mux.Close() 203 | }) 204 | taga, err := mux.Tag("a") 205 | if err != nil { 206 | skipIfProtocolNotSupported(t, err) 207 | assert.Nil(t, err) 208 | } 209 | tagb, _ := mux.Tag("b") 210 | tagc, _ := mux.Tag("c") 211 | assert.Nil(t, err) 212 | 213 | td, err := mux.ReadWhile(func() error { 214 | io.WriteString(taga, "out1") 215 | time.Sleep(sleepDuration) 216 | io.WriteString(tagb, "err1") 217 | io.WriteString(tagb, "err2") 218 | time.Sleep(sleepDuration) 219 | io.WriteString(tagc, "other") 220 | return nil 221 | }) 222 | 223 | assert.Equal(t, 3, len(td)) 224 | out1 := td[0] 225 | assert.Equal(t, "a", out1.Tag) 226 | assert.Equal(t, "out1", string(out1.Data)) 227 | err1 := td[1] 228 | assert.Equal(t, "b", err1.Tag) 229 | assert.Equal(t, "err1err2", string(err1.Data)) 230 | out2 := td[2] 231 | assert.Equal(t, "c", out2.Tag) 232 | assert.Equal(t, "other", string(out2.Data)) 233 | }) 234 | } 235 | } 236 | 237 | func TestMuxCmd(t *testing.T) { 238 | for _, network := range networks { 239 | t.Run(network, func(t *testing.T) { 240 | mux := &Mux[int]{network: network} 241 | t.Cleanup(func() { 242 | mux.Close() 243 | }) 244 | // sleep to avoid racing on connection oriented networks 245 | echo := fmt.Sprintf("echo out1 && sleep %f && echo err1 1>&2 && sleep %f && echo out2", sleepSecs, sleepSecs) 246 | cmd := exec.Command("sh", "-c", echo) 247 | stdout, err := mux.Tag(0) 248 | if err != nil { 249 | skipIfProtocolNotSupported(t, err) 250 | assert.Nil(t, err) 251 | } 252 | stderr, _ := mux.Tag(1) 253 | cmd.Stdout = stdout 254 | cmd.Stderr = stderr 255 | td, err := mux.ReadWhile(func() error { 256 | err := cmd.Run() 257 | return err 258 | }) 259 | 260 | assert.Equal(t, 3, len(td)) 261 | out1 := td[0] 262 | assert.Equal(t, 0, out1.Tag) 263 | assert.Equal(t, "out1\n", string(out1.Data)) 264 | err1 := td[1] 265 | assert.Equal(t, 1, err1.Tag) 266 | assert.Equal(t, "err1\n", string(err1.Data)) 267 | out2 := td[2] 268 | assert.Equal(t, 0, out2.Tag) 269 | assert.Equal(t, "out2\n", string(out2.Data)) 270 | }) 271 | } 272 | } 273 | -------------------------------------------------------------------------------- /iomux.go: -------------------------------------------------------------------------------- 1 | package iomux 2 | 3 | import ( 4 | "context" 5 | "errors" 6 | "fmt" 7 | "io" 8 | "net" 9 | "os" 10 | "path/filepath" 11 | "runtime" 12 | "sync" 13 | "time" 14 | ) 15 | 16 | // Mux provides a single receive and multiple send ends using unix domain networking. 17 | type Mux[T comparable] struct { 18 | network string 19 | dir string 20 | recvonce sync.Once 21 | recvaddr *net.UnixAddr 22 | recvconns []*net.UnixConn 23 | recvbufs [][]byte 24 | recvchan chan *taggedData[T] 25 | recvmutex []sync.Mutex 26 | recvstate map[recvKey]*recvState 27 | acceptFn func() error 28 | senders map[T]*net.UnixConn 29 | closed bool 30 | closers []io.Closer 31 | } 32 | 33 | type TaggedData[T comparable] struct { 34 | Tag T 35 | Data []byte 36 | } 37 | 38 | type taggedData[T comparable] struct { 39 | tag T 40 | data []byte 41 | conn *net.UnixConn 42 | err error 43 | } 44 | 45 | type recvKey struct { 46 | ctx context.Context 47 | conn *net.UnixConn 48 | } 49 | 50 | type recvState struct { 51 | eof bool 52 | } 53 | 54 | var ( 55 | MuxClosed = errors.New("mux has been closed") 56 | MuxNoConnections = errors.New("no senders have been connected") 57 | ) 58 | 59 | const deadlineDuration = 100 * time.Millisecond 60 | 61 | // Option to override defaults settings 62 | type Option[T comparable] func(*Mux[T]) 63 | 64 | func WithCustomDir[T comparable](dir string) Option[T] { 65 | return func(m *Mux[T]) { 66 | m.dir = dir 67 | } 68 | } 69 | 70 | // NewMuxUnix Create a new Mux using 'unix' network. 71 | func NewMuxUnix[T comparable]() *Mux[T] { 72 | return &Mux[T]{network: "unix"} 73 | } 74 | 75 | // NewMuxUnixGram Create a new Mux using 'unixgram' network. 76 | func NewMuxUnixGram[T comparable]() *Mux[T] { 77 | return &Mux[T]{network: "unixgram"} 78 | } 79 | 80 | // NewMuxUnixPacket Create a new Mux using 'unixpacket' network. 81 | func NewMuxUnixPacket[T comparable]() *Mux[T] { 82 | return &Mux[T]{network: "unixpacket"} 83 | } 84 | 85 | // Tag Create a file to receive data tagged with tag T. Returns an *os.File ready for writing, or an error. If an error 86 | // occurs when creating the receive end of the connection, the Mux will be closed. 87 | func (mux *Mux[T]) Tag(tag T, opts ...Option[T]) (*os.File, error) { 88 | if mux.closed { 89 | return nil, MuxClosed 90 | } 91 | err := mux.createReceiver() 92 | if err != nil { 93 | mux.Close() 94 | return nil, err 95 | } 96 | 97 | for _, opt := range opts { 98 | opt(mux) 99 | } 100 | 101 | sender, err := mux.createSender(tag) 102 | if err != nil { 103 | return nil, err 104 | } 105 | return sender, nil 106 | } 107 | 108 | // Read perform a read, blocking until data is available or ctx.Done. For connection oriented networks, Read 109 | // concurrently reads all connections buffering in order received for consecutive calls to Read. Returns io.EOF error 110 | // when there is no data remaining to be read. 111 | func (mux *Mux[T]) Read(ctx context.Context) ([]byte, T, error) { 112 | var zeroTag T 113 | if mux.closed { 114 | return nil, zeroTag, MuxClosed 115 | } 116 | if len(mux.recvconns) == 0 { 117 | return nil, zeroTag, MuxNoConnections 118 | } 119 | 120 | if len(mux.recvconns) == 1 { 121 | return mux.read(ctx, mux.recvconns[0], mux.recvbufs[0]) 122 | } 123 | 124 | if mux.recvstate == nil { 125 | mux.recvstate = make(map[recvKey]*recvState) 126 | } 127 | for i, c := range mux.recvconns { 128 | key := recvKey{ctx: ctx, conn: c} 129 | if _, ok := mux.recvstate[key]; !ok { 130 | mux.recvstate[key] = &recvState{} 131 | } else if mux.recvstate[key].eof { 132 | // avoid spinning up another read, we're done 133 | continue 134 | } 135 | conn := c 136 | buf := mux.recvbufs[i] 137 | mutex := &mux.recvmutex[i] 138 | if mutex.TryLock() { 139 | go func() { 140 | data, tag, err := mux.read(ctx, conn, buf) 141 | mux.recvchan <- &taggedData[T]{ 142 | data: data, 143 | tag: tag, 144 | conn: conn, 145 | err: err, 146 | } 147 | mutex.Unlock() 148 | }() 149 | } 150 | } 151 | 152 | sleepDuration := 1 * time.Millisecond 153 | for { 154 | select { 155 | case td := <-mux.recvchan: 156 | if td.err != nil { 157 | if td.err == io.EOF { 158 | key := recvKey{ctx: ctx, conn: td.conn} 159 | mux.recvstate[key].eof = true 160 | continue 161 | } 162 | return nil, zeroTag, td.err 163 | } 164 | return td.data, td.tag, td.err 165 | case <-ctx.Done(): 166 | done := true 167 | for _, c := range mux.recvconns { 168 | key := recvKey{ctx: ctx, conn: c} 169 | if state, ok := mux.recvstate[key]; ok { 170 | if !state.eof { 171 | done = false 172 | } 173 | } else { 174 | panic("no state") 175 | } 176 | } 177 | if done { 178 | for _, c := range mux.recvconns { 179 | key := recvKey{ctx: ctx, conn: c} 180 | delete(mux.recvstate, key) 181 | } 182 | return nil, zeroTag, io.EOF 183 | } 184 | default: 185 | } 186 | 187 | time.Sleep(sleepDuration) 188 | sleepDuration += sleepDuration 189 | if sleepDuration > deadlineDuration { 190 | sleepDuration = deadlineDuration 191 | } 192 | } 193 | } 194 | 195 | func (mux *Mux[T]) read(ctx context.Context, conn *net.UnixConn, buf []byte) ([]byte, T, error) { 196 | var zeroTag T 197 | for { 198 | _ = conn.SetDeadline(time.Now().Add(deadlineDuration)) 199 | n, addr, err := conn.ReadFrom(buf) 200 | if err != nil { 201 | if errors.Unwrap(err) != os.ErrDeadlineExceeded { 202 | return nil, zeroTag, err 203 | } 204 | select { 205 | case <-ctx.Done(): 206 | return nil, zeroTag, io.EOF 207 | default: 208 | continue 209 | } 210 | } 211 | data := make([]byte, n) 212 | copy(data, buf[0:n]) 213 | var tag T 214 | for t, c := range mux.senders { 215 | localAddr := c.LocalAddr().String() 216 | remoteAddr := conn.RemoteAddr() 217 | if addr != nil { 218 | if addr.String() == localAddr { 219 | tag = t 220 | break 221 | } 222 | } else if remoteAddr != nil && remoteAddr.String() == localAddr { 223 | tag = t 224 | break 225 | } 226 | } 227 | return data, tag, nil 228 | } 229 | } 230 | 231 | // ReadWhile Read until waitFn returns, returning the read data. 232 | func (mux *Mux[T]) ReadWhile(waitFn func() error) ([]*TaggedData[T], error) { 233 | if mux.closed { 234 | return nil, MuxClosed 235 | } 236 | ctx, cancelFn := context.WithCancel(context.Background()) 237 | var waitErr error 238 | go func() { 239 | waitErr = waitFn() 240 | cancelFn() 241 | }() 242 | td, err := mux.ReadUntil(ctx) 243 | if err != nil { 244 | return nil, err 245 | } 246 | return td, waitErr 247 | } 248 | 249 | // ReadUntil Read the receiver until done receives true 250 | func (mux *Mux[T]) ReadUntil(ctx context.Context) ([]*TaggedData[T], error) { 251 | if mux.closed { 252 | return nil, MuxClosed 253 | } 254 | var result []*TaggedData[T] 255 | for { 256 | data, tag, err := mux.Read(ctx) 257 | if err != nil { 258 | if err == io.EOF { 259 | return result, nil 260 | } 261 | return nil, err 262 | } 263 | resultLen := len(result) 264 | if resultLen > 0 { 265 | previous := result[resultLen-1] 266 | if previous.Tag == tag { 267 | previous.Data = append(previous.Data, data...) 268 | continue 269 | } 270 | } 271 | result = append(result, &TaggedData[T]{ 272 | Data: data, 273 | Tag: tag, 274 | }) 275 | } 276 | } 277 | 278 | // Close closes the Mux, closing connections and removing temporary files. Prevents reuse. 279 | func (mux *Mux[T]) Close() error { 280 | if mux.closed { 281 | return MuxClosed 282 | } 283 | mux.closed = true 284 | for _, closer := range mux.closers { 285 | closer.Close() 286 | } 287 | os.RemoveAll(mux.dir) 288 | return nil 289 | } 290 | 291 | func (mux *Mux[T]) createReceiver() (e error) { 292 | mux.recvonce.Do(func() { 293 | if mux.network == "" { 294 | switch runtime.GOOS { 295 | case "darwin": 296 | mux.network = "unix" 297 | default: 298 | mux.network = "unixgram" 299 | } 300 | } 301 | 302 | if mux.dir == "" { 303 | mux.dir, e = os.MkdirTemp("", "mux") 304 | if e != nil { 305 | return 306 | } 307 | } 308 | 309 | file := filepath.Join(mux.dir, "recv.sock") 310 | mux.recvaddr, e = net.ResolveUnixAddr(mux.network, file) 311 | if e != nil { 312 | return 313 | } 314 | mux.recvchan = make(chan *taggedData[T], 10) 315 | e = mux.startListener() 316 | if e != nil { 317 | return 318 | } 319 | }) 320 | return 321 | } 322 | 323 | func (mux *Mux[T]) startListener() error { 324 | // If we got at the underlying poll.FD it would be possible to call recvfrom with MSG_PEEK | MSG_TRUNC to size 325 | // the buffer to the current packet, but for now we just set the maximum message size for the OS for message 326 | // oriented unixgram, because the message truncates if it exceeds the buffer, and a modest read buffer otherwise. 327 | bufsize := 0 328 | switch mux.network { 329 | case "unixgram": 330 | { 331 | switch runtime.GOOS { 332 | case "darwin": 333 | bufsize = 2048 334 | default: 335 | bufsize = 65536 336 | } 337 | conn, err := net.ListenUnixgram(mux.network, mux.recvaddr) 338 | if err != nil { 339 | return err 340 | } 341 | mux.closers = append(mux.closers, conn) 342 | mux.acceptFn = func() error { 343 | return nil 344 | } 345 | _ = conn.CloseWrite() 346 | mux.recvconns = append(mux.recvconns, conn) 347 | mux.recvbufs = append(mux.recvbufs, make([]byte, bufsize)) 348 | mux.recvmutex = append(mux.recvmutex, sync.Mutex{}) 349 | } 350 | case "unix", "unixpacket": 351 | { 352 | bufsize = 256 353 | listener, err := net.ListenUnix(mux.network, mux.recvaddr) 354 | if err != nil { 355 | return err 356 | } 357 | mux.closers = append(mux.closers, listener) 358 | mux.acceptFn = func() error { 359 | err := listener.SetDeadline(time.Now().Add(deadlineDuration)) 360 | if err != nil { 361 | return err 362 | } 363 | conn, err := listener.AcceptUnix() 364 | if err != nil { 365 | return err 366 | } 367 | mux.closers = append(mux.closers, conn) 368 | _ = conn.CloseWrite() 369 | mux.recvconns = append(mux.recvconns, conn) 370 | mux.recvbufs = append(mux.recvbufs, make([]byte, bufsize)) 371 | mux.recvmutex = append(mux.recvmutex, sync.Mutex{}) 372 | return nil 373 | } 374 | } 375 | } 376 | 377 | return nil 378 | } 379 | 380 | func (mux *Mux[T]) createSender(tag T) (*os.File, error) { 381 | if mux.senders == nil { 382 | mux.senders = make(map[T]*net.UnixConn) 383 | } 384 | num := len(mux.senders) + 1 385 | if _, ok := mux.senders[tag]; !ok { 386 | address := filepath.Join(mux.dir, fmt.Sprintf("send_%d.sock", num)) 387 | addr, err := net.ResolveUnixAddr(mux.network, address) 388 | if err != nil { 389 | return nil, err 390 | } 391 | wg := sync.WaitGroup{} 392 | wg.Add(1) 393 | var acceptErr error 394 | go func() { 395 | acceptErr = mux.acceptFn() 396 | wg.Done() 397 | }() 398 | conn, dialErr := net.DialUnix(mux.network, addr, mux.recvaddr) 399 | wg.Wait() 400 | if acceptErr != nil { 401 | return nil, acceptErr 402 | } 403 | if dialErr != nil { 404 | return nil, dialErr 405 | } 406 | mux.closers = append(mux.closers, conn) 407 | _ = conn.CloseRead() 408 | mux.senders[tag] = conn 409 | } 410 | 411 | file, err := mux.senders[tag].File() 412 | if err != nil { 413 | return nil, err 414 | } 415 | return file, nil 416 | } 417 | -------------------------------------------------------------------------------- /LICENSE.md: -------------------------------------------------------------------------------- 1 | Apache License 2 | Version 2.0, January 2004 3 | http://www.apache.org/licenses/ 4 | 5 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 6 | 7 | 1. Definitions. 8 | 9 | "License" shall mean the terms and conditions for use, reproduction, 10 | and distribution as defined by Sections 1 through 9 of this document. 11 | 12 | "Licensor" shall mean the copyright owner or entity authorized by 13 | the copyright owner that is granting the License. 14 | 15 | "Legal Entity" shall mean the union of the acting entity and all 16 | other entities that control, are controlled by, or are under common 17 | control with that entity. For the purposes of this definition, 18 | "control" means (i) the power, direct or indirect, to cause the 19 | direction or management of such entity, whether by contract or 20 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 21 | outstanding shares, or (iii) beneficial ownership of such entity. 22 | 23 | "You" (or "Your") shall mean an individual or Legal Entity 24 | exercising permissions granted by this License. 25 | 26 | "Source" form shall mean the preferred form for making modifications, 27 | including but not limited to software source code, documentation 28 | source, and configuration files. 29 | 30 | "Object" form shall mean any form resulting from mechanical 31 | transformation or translation of a Source form, including but 32 | not limited to compiled object code, generated documentation, 33 | and conversions to other media types. 34 | 35 | "Work" shall mean the work of authorship, whether in Source or 36 | Object form, made available under the License, as indicated by a 37 | copyright notice that is included in or attached to the work 38 | (an example is provided in the Appendix below). 39 | 40 | "Derivative Works" shall mean any work, whether in Source or Object 41 | form, that is based on (or derived from) the Work and for which the 42 | editorial revisions, annotations, elaborations, or other modifications 43 | represent, as a whole, an original work of authorship. For the purposes 44 | of this License, Derivative Works shall not include works that remain 45 | separable from, or merely link (or bind by name) to the interfaces of, 46 | the Work and Derivative Works thereof. 47 | 48 | "Contribution" shall mean any work of authorship, including 49 | the original version of the Work and any modifications or additions 50 | to that Work or Derivative Works thereof, that is intentionally 51 | submitted to Licensor for inclusion in the Work by the copyright owner 52 | or by an individual or Legal Entity authorized to submit on behalf of 53 | the copyright owner. For the purposes of this definition, "submitted" 54 | means any form of electronic, verbal, or written communication sent 55 | to the Licensor or its representatives, including but not limited to 56 | communication on electronic mailing lists, source code control systems, 57 | and issue tracking systems that are managed by, or on behalf of, the 58 | Licensor for the purpose of discussing and improving the Work, but 59 | excluding communication that is conspicuously marked or otherwise 60 | designated in writing by the copyright owner as "Not a Contribution." 61 | 62 | "Contributor" shall mean Licensor and any individual or Legal Entity 63 | on behalf of whom a Contribution has been received by Licensor and 64 | subsequently incorporated within the Work. 65 | 66 | 2. Grant of Copyright License. Subject to the terms and conditions of 67 | this License, each Contributor hereby grants to You a perpetual, 68 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 69 | copyright license to reproduce, prepare Derivative Works of, 70 | publicly display, publicly perform, sublicense, and distribute the 71 | Work and such Derivative Works in Source or Object form. 72 | 73 | 3. Grant of Patent License. Subject to the terms and conditions of 74 | this License, each Contributor hereby grants to You a perpetual, 75 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 76 | (except as stated in this section) patent license to make, have made, 77 | use, offer to sell, sell, import, and otherwise transfer the Work, 78 | where such license applies only to those patent claims licensable 79 | by such Contributor that are necessarily infringed by their 80 | Contribution(s) alone or by combination of their Contribution(s) 81 | with the Work to which such Contribution(s) was submitted. If You 82 | institute patent litigation against any entity (including a 83 | cross-claim or counterclaim in a lawsuit) alleging that the Work 84 | or a Contribution incorporated within the Work constitutes direct 85 | or contributory patent infringement, then any patent licenses 86 | granted to You under this License for that Work shall terminate 87 | as of the date such litigation is filed. 88 | 89 | 4. Redistribution. You may reproduce and distribute copies of the 90 | Work or Derivative Works thereof in any medium, with or without 91 | modifications, and in Source or Object form, provided that You 92 | meet the following conditions: 93 | 94 | (a) You must give any other recipients of the Work or 95 | Derivative Works a copy of this License; and 96 | 97 | (b) You must cause any modified files to carry prominent notices 98 | stating that You changed the files; and 99 | 100 | (c) You must retain, in the Source form of any Derivative Works 101 | that You distribute, all copyright, patent, trademark, and 102 | attribution notices from the Source form of the Work, 103 | excluding those notices that do not pertain to any part of 104 | the Derivative Works; and 105 | 106 | (d) If the Work includes a "NOTICE" text file as part of its 107 | distribution, then any Derivative Works that You distribute must 108 | include a readable copy of the attribution notices contained 109 | within such NOTICE file, excluding those notices that do not 110 | pertain to any part of the Derivative Works, in at least one 111 | of the following places: within a NOTICE text file distributed 112 | as part of the Derivative Works; within the Source form or 113 | documentation, if provided along with the Derivative Works; or, 114 | within a display generated by the Derivative Works, if and 115 | wherever such third-party notices normally appear. The contents 116 | of the NOTICE file are for informational purposes only and 117 | do not modify the License. You may add Your own attribution 118 | notices within Derivative Works that You distribute, alongside 119 | or as an addendum to the NOTICE text from the Work, provided 120 | that such additional attribution notices cannot be construed 121 | as modifying the License. 122 | 123 | You may add Your own copyright statement to Your modifications and 124 | may provide additional or different license terms and conditions 125 | for use, reproduction, or distribution of Your modifications, or 126 | for any such Derivative Works as a whole, provided Your use, 127 | reproduction, and distribution of the Work otherwise complies with 128 | the conditions stated in this License. 129 | 130 | 5. Submission of Contributions. Unless You explicitly state otherwise, 131 | any Contribution intentionally submitted for inclusion in the Work 132 | by You to the Licensor shall be under the terms and conditions of 133 | this License, without any additional terms or conditions. 134 | Notwithstanding the above, nothing herein shall supersede or modify 135 | the terms of any separate license agreement you may have executed 136 | with Licensor regarding such Contributions. 137 | 138 | 6. Trademarks. This License does not grant permission to use the trade 139 | names, trademarks, service marks, or product names of the Licensor, 140 | except as required for reasonable and customary use in describing the 141 | origin of the Work and reproducing the content of the NOTICE file. 142 | 143 | 7. Disclaimer of Warranty. Unless required by applicable law or 144 | agreed to in writing, Licensor provides the Work (and each 145 | Contributor provides its Contributions) on an "AS IS" BASIS, 146 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 147 | implied, including, without limitation, any warranties or conditions 148 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 149 | PARTICULAR PURPOSE. You are solely responsible for determining the 150 | appropriateness of using or redistributing the Work and assume any 151 | risks associated with Your exercise of permissions under this License. 152 | 153 | 8. Limitation of Liability. In no event and under no legal theory, 154 | whether in tort (including negligence), contract, or otherwise, 155 | unless required by applicable law (such as deliberate and grossly 156 | negligent acts) or agreed to in writing, shall any Contributor be 157 | liable to You for damages, including any direct, indirect, special, 158 | incidental, or consequential damages of any character arising as a 159 | result of this License or out of the use or inability to use the 160 | Work (including but not limited to damages for loss of goodwill, 161 | work stoppage, computer failure or malfunction, or any and all 162 | other commercial damages or losses), even if such Contributor 163 | has been advised of the possibility of such damages. 164 | 165 | 9. Accepting Warranty or Additional Liability. While redistributing 166 | the Work or Derivative Works thereof, You may choose to offer, 167 | and charge a fee for, acceptance of support, warranty, indemnity, 168 | or other liability obligations and/or rights consistent with this 169 | License. However, in accepting such obligations, You may act only 170 | on Your own behalf and on Your sole responsibility, not on behalf 171 | of any other Contributor, and only if You agree to indemnify, 172 | defend, and hold each Contributor harmless for any liability 173 | incurred by, or claims asserted against, such Contributor by reason 174 | of your accepting any such warranty or additional liability. 175 | 176 | END OF TERMS AND CONDITIONS 177 | 178 | APPENDIX: How to apply the Apache License to your work. 179 | 180 | To apply the Apache License to your work, attach the following 181 | boilerplate notice, with the fields enclosed by brackets "[]" 182 | replaced with your own identifying information. (Don't include 183 | the brackets!) The text should be enclosed in the appropriate 184 | comment syntax for the file format. We also recommend that a 185 | file or class name and description of purpose be included on the 186 | same "printed page" as the copyright notice for easier 187 | identification within third-party archives. 188 | 189 | Copyright [yyyy] [name of copyright owner] 190 | 191 | Licensed under the Apache License, Version 2.0 (the "License"); 192 | you may not use this file except in compliance with the License. 193 | You may obtain a copy of the License at 194 | 195 | http://www.apache.org/licenses/LICENSE-2.0 196 | 197 | Unless required by applicable law or agreed to in writing, software 198 | distributed under the License is distributed on an "AS IS" BASIS, 199 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 200 | See the License for the specific language governing permissions and 201 | limitations under the License. 202 | --------------------------------------------------------------------------------