github.com/bingoohuang/gg@v0.0.0-20240325092523-45da7dee9335/pkg/iox/maxlatencywriter.go (about)

     1  package iox
     2  
     3  import (
     4  	"io"
     5  	"log"
     6  	"sync"
     7  	"time"
     8  )
     9  
    10  // from https://golang.org/src/net/http/httputil/reverseproxy.go
    11  
    12  type WriteFlusher interface {
    13  	io.Writer
    14  	Flush() error
    15  }
    16  
    17  type MaxLatencyWriter struct {
    18  	Dst     WriteFlusher
    19  	Latency time.Duration // non-zero; negative means to flush immediately
    20  
    21  	mu           sync.Mutex // protects t, flushPending, and dst.Flush
    22  	t            *time.Timer
    23  	flushPending bool
    24  }
    25  
    26  func (m *MaxLatencyWriter) Write(p []byte) (n int, err error) {
    27  	m.mu.Lock()
    28  	defer m.mu.Unlock()
    29  	n, err = m.Dst.Write(p)
    30  	if m.Latency < 0 {
    31  		if err := m.Dst.Flush(); err != nil {
    32  			return 0, err
    33  		}
    34  		return
    35  	}
    36  	if m.flushPending {
    37  		return
    38  	}
    39  	if m.t == nil {
    40  		m.t = time.AfterFunc(m.Latency, m.delayedFlush)
    41  	} else {
    42  		m.t.Reset(m.Latency)
    43  	}
    44  	m.flushPending = true
    45  	return
    46  }
    47  
    48  func (m *MaxLatencyWriter) delayedFlush() {
    49  	m.mu.Lock()
    50  	defer m.mu.Unlock()
    51  	if !m.flushPending { // if stop was called but AfterFunc already started this goroutine
    52  		return
    53  	}
    54  	if err := m.Dst.Flush(); err != nil {
    55  		log.Printf("flush failed: %v", err)
    56  	}
    57  	m.flushPending = false
    58  }
    59  
    60  func (m *MaxLatencyWriter) Stop() {
    61  	m.mu.Lock()
    62  	defer m.mu.Unlock()
    63  	m.flushPending = false
    64  	if m.t != nil {
    65  		m.t.Stop()
    66  	}
    67  }