lib9p

client.go (11236B)

---

 package lib9p
 
 import (
 	"context"
 	"fmt"
 	"io"
 	"sync"
 )
 
 type Client struct {
 	msize     uint32
 	mSizeLock *sync.Mutex
 	uname     string
 	fPool     *clientFidPool
 	txc       chan<- *clientReq
 	errc      chan error
 	cancel    context.CancelFunc
 	rootFid   *clientFid
 	wg        *sync.WaitGroup
 }
 
 func NewClient(mSize uint32, uname string, r io.Reader, w io.Writer) *Client {
 	ctx, cancel := context.WithCancel(context.Background())
 	c := &Client{
 		msize:     mSize,
 		mSizeLock: new(sync.Mutex),
 		uname:     uname,
 		fPool:     allocClientFidPool(),
 		errc:      make(chan error),
 		cancel:    cancel,
 		wg:        new(sync.WaitGroup),
 	}
 	tmsgc := c.runSpeaker(ctx, w)
 	rmsgc := c.runListener(ctx, r)
 	c.txc = c.runMultiplexer(ctx, tmsgc, rmsgc)
 	return c
 }
 
 func (c *Client) Stop() {
 	c.cancel()
 	c.wg.Wait()
 	close(c.errc)
 }
 
 func (c *Client) mSize() uint32 {
 	c.mSizeLock.Lock()
 	defer c.mSizeLock.Unlock()
 	return c.msize
 }
 
 func (c *Client) setMSize(mSize uint32) {
 	c.mSizeLock.Lock()
 	defer c.mSizeLock.Unlock()
 	c.msize = mSize
 }
 
 // RunListener runs listener goroutine.
 // Listener reads byte array of 9P messages from r and make each of them into
 // corresponding struct that implements Msg, and sends it to the returned channel.
 // Listener goroutine returns when ctx is canceled.
 // Listener goroutine reports errors to the client's errc channel.
 func (c *Client) runListener(ctx context.Context, r io.Reader) <-chan Msg {
 	c.wg.Add(1)
 	// TODO: terminate with ctx.Done()
 	rmsgc := make(chan Msg, 3)
 	go func() {
 		wg := new(sync.WaitGroup)
 		defer func() {
 			wg.Wait()
 			close(rmsgc)
 			c.wg.Done()
 		}()
 		for {
 			select {
 			case <-ctx.Done():
 				// TODO: should return error via ec??
 				// TODO: should close r?
 				return
 			default:
 				done := make(chan struct{})
 				var (
 					msg Msg
 					err error
 				)
 				go func() {
 					defer close(done)
 					msg, err = recv(r)
 				}()
 				select {
 				case <-done:
 				case <-ctx.Done():
 				}
 				if err != nil {
 					c.errc <- fmt.Errorf("recv: %v", err)
 					continue
 				}
 				wg.Add(1)
 				go func() {
 					defer wg.Done()
 					select {
 					case rmsgc <- msg:
 					case <-ctx.Done():
 					}
 				}()
 			}
 		}
 	}()
 	return rmsgc
 }
 
 // RunSpeaker runs speaker goroutine.
 // Speaker goroutine recieves 9P Msgs from the returned channel, marshal them
 // into byte arrays and sends them to w.
 // It reports any errors to the clients errc channel.
 // It returnes when ctx is canceled.
 func (c *Client) runSpeaker(ctx context.Context, w io.Writer) chan<- Msg {
 	c.wg.Add(1)
 	tmsgc := make(chan Msg, 3)
 	go func() {
 		defer c.wg.Done()
 		for {
 			select {
 			case <-ctx.Done():
 				return
 			case msg := <-tmsgc:
 				if msg == nil {
 					// tmsgc is closed, which means ctx.Done() is also closed.
 					// but this code breaks semantics?
 					return
 				}
 				if err := send(msg, w); err != nil {
 					c.errc <- fmt.Errorf("send: %v", err)
 				}
 			}
 		}
 	}()
 	return tmsgc
 }
 
 // RunMultiplexer runs multiplexer goroutine.
 // Multiplexer goroutines, one for recieving Rmsg and another for sending Tmsg.
 // The goroutine for Tmsg recieves *clientReq from the returned channel,
 // and send the 9P Msg to the speaker goroutine via tmsgc.
 // The goroutine for Rmsg recieves *clientReq from the Tmsg goroutine and waits for
 // the reply to the corresponding message from the listener goroutine via rmsgc.
 // After recieving the reply, it sets the *clientReq.rmsg and sends it t the
 // *clientReq.rxc.
 // It reports any errors to the client's errc channel.
 func (c *Client) runMultiplexer(ctx context.Context, tmsgc chan<- Msg, rmsgc <-chan Msg) chan<- *clientReq {
 	c.wg.Add(2)
 	txc := make(chan *clientReq)
 	reqc := make(chan *clientReq)
 	// Rmsg
 	go func(reqc <-chan *clientReq) {
 		wg := new(sync.WaitGroup)
 		defer func() {
 			wg.Wait()
 			c.wg.Done()
 		}()
 		rPool := make(map[uint16]*clientReq)
 		for {
 			select {
 			case <-ctx.Done():
 				return
 			case req := <-reqc:
 				if req == nil {
 					// ctx is canceled.
 					continue
 				}
 				if _, ok := rPool[req.tag]; ok {
 					c.errc <- fmt.Errorf("mux: duplicate tag: %d", req.tag)
 					continue
 				}
 				rPool[req.tag] = req // TODO: wait for req.ctxDone channel.
 				wg.Add(1)
 				go func() {
 					defer wg.Done()
 					<-req.ctxDone
 				}()
 			case msg := <-rmsgc:
 				if msg == nil {
 					// ctx is canceled.
 					continue
 				}
 				req, ok := rPool[msg.Tag()]
 				if !ok {
 					c.errc <- fmt.Errorf("mux: unknown tag for msg: %v", msg)
 					continue
 				}
 				delete(rPool, msg.Tag())
 				req.rmsg = msg
 				go func() {
 					defer close(req.rxc)
 					select {
 					case <-req.ctxDone:
 					case req.rxc <- req:
 					}
 				}()
 			}
 		}
 	}(reqc)
 	// Tmsg
 	go func(reqc chan<- *clientReq) {
 		wg := new(sync.WaitGroup)
 		defer func() {
 			wg.Wait()
 			close(reqc)
 			close(tmsgc)
 			c.wg.Done()
 		}()
 		for {
 			select {
 			case <-ctx.Done():
 				return
 			case req := <-txc:
 				select {
 				case reqc <- req:
 				case <-ctx.Done():
 					return
 				}
 				wg.Add(1)
 				go func() {
 					defer wg.Done()
 					tmsgc <- req.tmsg
 				}()
 			}
 		}
 	}(reqc)
 	return txc
 }
 
 // Transact send 9P Msg of req to the multiplexer goroutines and recieves
 // the reply.
 func (c *Client) transact(ctx context.Context, tmsg Msg) (Msg, error) {
 	ctx1, cancel1 := context.WithCancel(ctx)
 	req := newClientReq(ctx1, tmsg)
 	select {
 	case <-ctx.Done():
 		return nil, ctx.Err()
 	case c.txc <- req:
 	}
 	select {
 	case req = <-req.rxc: // TODO: this assignment is not required.
 		cancel1()
 		return req.rmsg, req.err
 	case <-ctx.Done():
 		return nil, ctx.Err()
 	}
 }
 
 func (c *Client) Version(ctx context.Context, tag uint16, mSize uint32, version string) (uint32, string, error) {
 	tmsg := &TVersion{tag: tag, mSize: mSize, version: version}
 	rmsg, err := c.transact(ctx, tmsg)
 	if err != nil {
 		return 0, "", fmt.Errorf("transact: %v", err)
 	}
 	switch rmsg := rmsg.(type) {
 	case *RVersion:
 		return rmsg.mSize, rmsg.version, nil
 	case *RError:
 		return 0, "", rmsg.ename
 	default:
 		return 0, "", fmt.Errorf("invalid reply: %v", rmsg)
 	}
 }
 
 func (c *Client) Auth(ctx context.Context, tag uint16, afid uint32, uname, aname string) (Qid, error) {
 	tmsg := &TAuth{tag: tag, afid: afid, uname: uname}
 	rmsg, err := c.transact(ctx, tmsg)
 	if err != nil {
 		return Qid{}, fmt.Errorf("transact: %v", err)
 	}
 	switch rmsg := rmsg.(type) {
 	case *RAuth:
 		return rmsg.aqid, nil
 	case *RError:
 		return Qid{}, rmsg.ename
 	default:
 		return Qid{}, fmt.Errorf("invalid reply: %v", rmsg)
 	}
 }
 
 func (c *Client) Attach(ctx context.Context, tag uint16, fid, afid uint32, uname, aname string) (Qid, error) {
 	tmsg := &TAttach{tag: tag, fid: fid, afid: afid, uname: uname, aname: aname}
 	rmsg, err := c.transact(ctx, tmsg)
 	if err != nil {
 		return Qid{}, fmt.Errorf("transact: %v", err)
 	}
 	switch rmsg := rmsg.(type) {
 	case *RAttach:
 		return rmsg.qid, nil
 	case *RError:
 		return Qid{}, rmsg.ename
 	default:
 		return Qid{}, fmt.Errorf("invalid reply: %v", rmsg)
 	}
 }
 
 func (c *Client) Flush(ctx context.Context, tag, oldtag uint16) error {
 	tmsg := &TFlush{tag: tag, oldtag: oldtag}
 	rmsg, err := c.transact(ctx, tmsg)
 	if err != nil {
 		return fmt.Errorf("transact: %v", err)
 	}
 	switch rmsg := rmsg.(type) {
 	case *RFlush:
 		return nil
 	case *RError:
 		return rmsg.ename
 	default:
 		return fmt.Errorf("invalid reply: %v", rmsg)
 	}
 }
 func (c *Client) Walk(ctx context.Context, tag uint16, fid, newFid uint32, wname []string) (wqid []Qid, err error) {
 	tmsg := &TWalk{tag: tag, fid: fid, newFid: newFid, wname: wname}
 	rmsg, err := c.transact(ctx, tmsg)
 	if err != nil {
 		return nil, fmt.Errorf("transact: %v", err)
 	}
 	switch rmsg := rmsg.(type) {
 	case *RWalk:
 		return rmsg.qid, nil
 	case *RError:
 		return nil, rmsg.ename
 	default:
 		return nil, fmt.Errorf("invalid reply: %v", rmsg)
 	}
 }
 func (c *Client) Open(ctx context.Context, tag uint16, fid uint32, mode OpenMode) (qid Qid, iounit uint32, err error) {
 	tmsg := &TOpen{tag: tag, fid: fid, mode: mode}
 	rmsg, err := c.transact(ctx, tmsg)
 	if err != nil {
 		return Qid{}, 0, fmt.Errorf("transact: %v", err)
 	}
 	switch rmsg := rmsg.(type) {
 	case *ROpen:
 		return rmsg.qid, rmsg.iounit, nil
 	case *RError:
 		return Qid{}, 0, rmsg.ename
 	default:
 		return Qid{}, 0, fmt.Errorf("invalid reply: %v", rmsg)
 	}
 }
 func (c *Client) Create(ctx context.Context, tag uint16, fid uint32, name string, perm FileMode, mode OpenMode) (Qid, uint32, error) {
 	tmsg := &TCreate{tag: tag, fid: fid, name: name, perm: perm, mode: mode}
 	rmsg, err := c.transact(ctx, tmsg)
 	if err != nil {
 		return Qid{}, 0, fmt.Errorf("transact: %v", err)
 	}
 	switch rmsg := rmsg.(type) {
 	case *RCreate:
 		return rmsg.qid, rmsg.iounit, nil
 	case *RError:
 		return Qid{}, 0, rmsg.ename
 	default:
 		return Qid{}, 0, fmt.Errorf("invalid reply: %v", rmsg)
 	}
 }
 func (c *Client) Read(ctx context.Context, tag uint16, fid uint32, offset uint64, count uint32) (data []byte, err error) {
 	tmsg := &TRead{tag: tag, fid: fid, offset: offset, count: count}
 	rmsg, err := c.transact(ctx, tmsg)
 	if err != nil {
 		return nil, fmt.Errorf("transact: %v", err)
 	}
 	switch rmsg := rmsg.(type) {
 	case *RRead:
 		return rmsg.data, nil
 	case *RError:
 		return nil, rmsg.ename
 	default:
 		return nil, fmt.Errorf("invalid reply: %v", rmsg)
 	}
 }
 func (c *Client) Write(ctx context.Context, tag uint16, fid uint32, offset uint64, count uint32, data []byte) (uint32, error) {
 	tmsg := &TWrite{tag: tag, fid: fid, offset: offset, count: count, data: data}
 	rmsg, err := c.transact(ctx, tmsg)
 	if err != nil {
 		return 0, fmt.Errorf("transact: %v", err)
 	}
 	switch rmsg := rmsg.(type) {
 	case *RWrite:
 		return rmsg.count, nil
 	case *RError:
 		return 0, rmsg.ename
 	default:
 		return 0, fmt.Errorf("invalid reply: %v", rmsg)
 	}
 }
 func (c *Client) Clunk(ctx context.Context, tag uint16, fid uint32) error {
 	tmsg := &TClunk{tag: tag, fid: fid}
 	rmsg, err := c.transact(ctx, tmsg)
 	if err != nil {
 		return fmt.Errorf("transact: %v", err)
 	}
 	switch rmsg := rmsg.(type) {
 	case *RClunk:
 		return nil
 	case *RError:
 		return rmsg.ename
 	default:
 		return fmt.Errorf("invalid reply: %v", rmsg)
 	}
 }
 func (c *Client) Remove(ctx context.Context, tag uint16, fid uint32) error {
 	tmsg := &TRemove{tag: tag, fid: fid}
 	rmsg, err := c.transact(ctx, tmsg)
 	if err != nil {
 		return fmt.Errorf("transact: %v", err)
 	}
 	switch rmsg := rmsg.(type) {
 	case *RRemove:
 		return nil
 	case *RError:
 		return rmsg.ename
 	default:
 		return fmt.Errorf("invalid reply: %v", rmsg)
 	}
 }
 func (c *Client) Stat(ctx context.Context, tag uint16, fid uint32) (*Stat, error) {
 	tmsg := &TStat{tag: tag, fid: fid}
 	rmsg, err := c.transact(ctx, tmsg)
 	if err != nil {
 		return nil, fmt.Errorf("transact: %v", err)
 	}
 	switch rmsg := rmsg.(type) {
 	case *RStat:
 		return rmsg.stat, nil
 	case *RError:
 		return nil, rmsg.ename
 	default:
 		return nil, fmt.Errorf("invalid reply: %v", rmsg)
 	}
 }
 func (c *Client) Wstat(ctx context.Context, tag uint16, fid uint32, stat *Stat) error {
 	tmsg := &TWStat{tag: tag, fid: fid, stat: stat}
 	rmsg, err := c.transact(ctx, tmsg)
 	if err != nil {
 		return fmt.Errorf("transact: %v", err)
 	}
 	switch rmsg := rmsg.(type) {
 	case *RWStat:
 		return nil
 	case *RError:
 		return rmsg.ename
 	default:
 		return fmt.Errorf("invalid reply: %v", rmsg)
 	}
 }