github.com/Asutorufa/yuhaiin@v0.3.6-0.20240502055049-7984da7023a0/pkg/net/deadline/deadline.go (about)

     1  package deadline
     2  
     3  import (
     4  	"context"
     5  	"time"
     6  )
     7  
     8  type Deadline struct {
     9  	ctx    context.Context
    10  	cancel context.CancelFunc
    11  
    12  	deadline *time.Timer
    13  
    14  	close func()
    15  }
    16  
    17  type opts struct {
    18  	close func()
    19  
    20  	wclose func()
    21  	rclose func()
    22  }
    23  
    24  func WithClose(f func()) func(*opts) {
    25  	return func(o *opts) {
    26  		o.close = f
    27  	}
    28  }
    29  
    30  func WithWriteClose(f func()) func(*opts) {
    31  	return func(o *opts) {
    32  		o.wclose = f
    33  	}
    34  }
    35  
    36  func WithReadClose(f func()) func(*opts) {
    37  	return func(o *opts) {
    38  		o.rclose = f
    39  	}
    40  }
    41  
    42  func New(os ...func(*opts)) *Deadline {
    43  	o := &opts{}
    44  	for _, f := range os {
    45  		f(o)
    46  	}
    47  
    48  	ctx, cancel := context.WithCancel(context.Background())
    49  
    50  	return &Deadline{
    51  		ctx:    ctx,
    52  		cancel: cancel,
    53  		close:  o.close,
    54  	}
    55  }
    56  
    57  func (x *Deadline) Close() error {
    58  	x.cancel()
    59  
    60  	deadline := x.deadline
    61  
    62  	if deadline != nil {
    63  		deadline.Stop()
    64  	}
    65  	return nil
    66  }
    67  
    68  func (x *Deadline) SetDeadline(t time.Time) {
    69  	until := time.Until(t)
    70  
    71  	if !t.IsZero() && until <= 0 {
    72  		_ = x.Close()
    73  		if x.close != nil {
    74  			x.close()
    75  		}
    76  		return
    77  	}
    78  
    79  	if x.deadline == nil {
    80  		if !t.IsZero() {
    81  			x.deadline = time.AfterFunc(until, func() {
    82  				x.cancel()
    83  				if x.close != nil {
    84  					x.close()
    85  				}
    86  			})
    87  		}
    88  		return
    89  	}
    90  
    91  	if t.IsZero() {
    92  		x.deadline.Stop()
    93  	} else {
    94  		x.deadline.Reset(until)
    95  	}
    96  }
    97  
    98  func (x *Deadline) Context() context.Context { return x.ctx }
    99  
   100  type PipeDeadline struct {
   101  	w *Deadline
   102  	r *Deadline
   103  }
   104  
   105  func NewPipe(os ...func(*opts)) *PipeDeadline {
   106  	o := &opts{}
   107  	for _, f := range os {
   108  		f(o)
   109  	}
   110  
   111  	w := New(WithClose(o.wclose))
   112  	r := New(WithClose(o.rclose))
   113  
   114  	return &PipeDeadline{w: w, r: r}
   115  }
   116  
   117  func (x *PipeDeadline) Close() error {
   118  	_ = x.r.Close()
   119  	return x.w.Close()
   120  }
   121  
   122  func (x *PipeDeadline) SetDeadline(t time.Time) {
   123  	x.w.SetDeadline(t)
   124  	x.r.SetDeadline(t)
   125  }
   126  
   127  func (x *PipeDeadline) SetWriteDeadline(t time.Time) {
   128  	x.w.SetDeadline(t)
   129  }
   130  
   131  func (x *PipeDeadline) SetReadDeadline(t time.Time) {
   132  	x.r.SetDeadline(t)
   133  }
   134  
   135  func (x *PipeDeadline) WriteContext() context.Context { return x.w.ctx }
   136  
   137  func (x *PipeDeadline) ReadContext() context.Context { return x.r.ctx }