github.com/cmd-stream/base-go@v0.0.0-20230813145615-dd6ac24c16f5/server/worker.go (about)

     1  package server
     2  
     3  import (
     4  	"context"
     5  	"net"
     6  
     7  	"github.com/cmd-stream/base-go"
     8  )
     9  
    10  // NewWorker creates a new Worker.
    11  func NewWorker(conns <-chan net.Conn, delegate base.ServerDelegate,
    12  	callback LostConnCallback,
    13  ) Worker {
    14  	ctx, cancel := context.WithCancel(context.Background())
    15  	return Worker{ctx, cancel, conns, delegate, callback}
    16  }
    17  
    18  // Worker is a Server worker.
    19  //
    20  // It receives a connections from the conns channel and handles them using the
    21  // delegate. Also it implements jointwork.Task, so it + ConnReceiver may do the
    22  // job together.
    23  //
    24  // If the connection processing failed with an error, it is passed to the
    25  // LostConnCallback.
    26  type Worker struct {
    27  	ctx      context.Context
    28  	cancel   context.CancelFunc
    29  	conns    <-chan net.Conn
    30  	delegate base.ServerDelegate
    31  	callback LostConnCallback
    32  }
    33  
    34  func (w Worker) Run() (err error) {
    35  	var (
    36  		conn net.Conn
    37  		more bool
    38  	)
    39  	for {
    40  		select {
    41  		case <-w.ctx.Done():
    42  			return ErrClosed
    43  		case conn, more = <-w.conns:
    44  			if !more {
    45  				return nil
    46  			}
    47  			err = w.delegate.Handle(w.ctx, conn)
    48  			if err != nil {
    49  				if err == context.Canceled {
    50  					err = ErrClosed
    51  				}
    52  				if w.callback != nil {
    53  					w.callback(conn.RemoteAddr(), err)
    54  				}
    55  				if err == ErrClosed {
    56  					return
    57  				}
    58  			}
    59  		}
    60  	}
    61  }
    62  
    63  func (w Worker) Stop() (err error) {
    64  	w.cancel()
    65  	return
    66  }