mnet: fix PacketConn support

This commit is contained in:
Brian Picciano 2019-03-01 14:21:43 -05:00
parent 14de888300
commit 5bd3bf1d6d

View File

@ -9,13 +9,18 @@ import (
"github.com/mediocregopher/mediocre-go-lib/mcfg" "github.com/mediocregopher/mediocre-go-lib/mcfg"
"github.com/mediocregopher/mediocre-go-lib/mctx" "github.com/mediocregopher/mediocre-go-lib/mctx"
"github.com/mediocregopher/mediocre-go-lib/merr"
"github.com/mediocregopher/mediocre-go-lib/mlog" "github.com/mediocregopher/mediocre-go-lib/mlog"
"github.com/mediocregopher/mediocre-go-lib/mrun" "github.com/mediocregopher/mediocre-go-lib/mrun"
) )
// Listener is returned by WithListen and simply wraps a net.Listener. // Listener is returned by WithListen and simply wraps a net.Listener.
type Listener struct { type Listener struct {
// One of these will be populated during the start hook, depending on the
// protocol configured.
net.Listener net.Listener
net.PacketConn
ctx context.Context ctx context.Context
// If set to true before mrun's stop event is run, the stop event will not // If set to true before mrun's stop event is run, the stop event will not
@ -28,6 +33,13 @@ type listenerOpts struct {
defaultAddr string defaultAddr string
} }
func (lOpts listenerOpts) isPacketConn() bool {
proto := strings.ToLower(lOpts.proto)
return strings.HasPrefix(proto, "udp") ||
proto == "unixgram" ||
strings.HasPrefix(proto, "ip")
}
// ListenerOpt is a value which adjusts the behavior of WithListener. // ListenerOpt is a value which adjusts the behavior of WithListener.
type ListenerOpt func(*listenerOpts) type ListenerOpt func(*listenerOpts)
@ -69,12 +81,21 @@ func WithListener(ctx context.Context, opts ...ListenerOpt) (context.Context, *L
l.ctx = mrun.WithStartHook(l.ctx, func(context.Context) error { l.ctx = mrun.WithStartHook(l.ctx, func(context.Context) error {
var err error var err error
if l.Listener, err = net.Listen(lOpts.proto, *addr); err != nil {
return err
}
l.ctx = mctx.Annotate(l.ctx, l.ctx = mctx.Annotate(l.ctx,
"proto", lOpts.proto, "proto", lOpts.proto,
"addr", l.Addr().String()) "addr", *addr)
if lOpts.isPacketConn() {
l.PacketConn, err = net.ListenPacket(lOpts.proto, *addr)
} else {
l.Listener, err = net.Listen(lOpts.proto, *addr)
}
if err != nil {
return merr.Wrap(err, l.ctx)
}
l.ctx = mctx.Annotate(l.ctx, "addr", l.Addr().String())
mlog.Info("listening", l.ctx) mlog.Info("listening", l.ctx)
return nil return nil
}) })