github.com/bluenviron/gomavlib/v2@v2.2.1-0.20240308101627-2c07e3da629c/pkg/reconnector/reconnector.go (about)

     1  // Package reconnector allows to perform automatic reconnections.
     2  package reconnector
     3  
     4  import (
     5  	"context"
     6  	"io"
     7  	"sync"
     8  	"time"
     9  )
    10  
    11  var reconnectPeriod = 2 * time.Second
    12  
    13  // ConnectFunc is the prototype of the callback passed to New()
    14  type ConnectFunc func(context.Context) (io.ReadWriteCloser, error)
    15  
    16  type connWithContext struct {
    17  	rwc       io.ReadWriteCloser
    18  	mutex     sync.Mutex
    19  	ctx       context.Context
    20  	ctxCancel func()
    21  }
    22  
    23  func newConnWithContext(rwc io.ReadWriteCloser) *connWithContext {
    24  	ctx, ctxCancel := context.WithCancel(context.Background())
    25  
    26  	return &connWithContext{
    27  		rwc:       rwc,
    28  		ctx:       ctx,
    29  		ctxCancel: ctxCancel,
    30  	}
    31  }
    32  
    33  func (c *connWithContext) Close() error {
    34  	c.mutex.Lock()
    35  	defer c.mutex.Unlock()
    36  
    37  	select {
    38  	case <-c.ctx.Done():
    39  		return nil
    40  	default:
    41  	}
    42  
    43  	c.ctxCancel()
    44  
    45  	return c.rwc.Close()
    46  }
    47  
    48  func (c *connWithContext) Read(p []byte) (int, error) {
    49  	n, err := c.rwc.Read(p)
    50  	if n == 0 {
    51  		c.Close() //nolint:errcheck
    52  	}
    53  	return n, err
    54  }
    55  
    56  func (c *connWithContext) Write(p []byte) (int, error) {
    57  	n, err := c.rwc.Write(p)
    58  	if n == 0 {
    59  		c.Close() //nolint:errcheck
    60  	}
    61  	return n, err
    62  }
    63  
    64  // Reconnector allows to perform automatic reconnections.
    65  type Reconnector struct {
    66  	connect ConnectFunc
    67  
    68  	ctx       context.Context
    69  	ctxCancel func()
    70  	curConn   *connWithContext
    71  }
    72  
    73  // New allocates a Reconnector.
    74  func New(connect ConnectFunc) *Reconnector {
    75  	ctx, ctxCancel := context.WithCancel(context.Background())
    76  
    77  	return &Reconnector{
    78  		connect:   connect,
    79  		ctx:       ctx,
    80  		ctxCancel: ctxCancel,
    81  	}
    82  }
    83  
    84  // Close closes a reconnector.
    85  func (a *Reconnector) Close() {
    86  	a.ctxCancel()
    87  }
    88  
    89  // Reconnect returns the next connection.
    90  func (a *Reconnector) Reconnect() (io.ReadWriteCloser, bool) {
    91  	if a.curConn != nil {
    92  		select {
    93  		case <-a.curConn.ctx.Done():
    94  		case <-a.ctx.Done():
    95  			return nil, false
    96  		}
    97  	}
    98  
    99  	for {
   100  		conn, err := a.connect(a.ctx)
   101  		if err != nil {
   102  			select {
   103  			case <-time.After(reconnectPeriod):
   104  				continue
   105  			case <-a.ctx.Done():
   106  				return nil, false
   107  			}
   108  		}
   109  
   110  		a.curConn = newConnWithContext(conn)
   111  		return a.curConn, true
   112  	}
   113  }