Add SyncWrite method to doppel.Conn

mtglib/internal/doppel/conn.go

@@ -16,18 +16,48 @@ type Conn struct {
 }
 
 type connPayload struct {
-	ctx         context.Context
-	ctxCancel   context.CancelCauseFunc
-	clock       Clock
-	wg          sync.WaitGroup
-	writeLock   sync.Mutex
-	writeStream bytes.Buffer
+	ctx           context.Context
+	ctxCancel     context.CancelCauseFunc
+	clock         Clock
+	wg            sync.WaitGroup
+	syncWriteLock sync.RWMutex
+	writeStream   bytes.Buffer
+	writeCond     *sync.Cond
 }
 
 func (c Conn) Write(p []byte) (int, error) {
-	c.p.writeLock.Lock()
+	c.p.syncWriteLock.RLock()
+	defer c.p.syncWriteLock.RUnlock()
+
+	c.p.writeCond.L.Lock()
 	c.p.writeStream.Write(p)
-	c.p.writeLock.Unlock()
+	c.p.writeCond.L.Unlock()
+
+	return len(p), context.Cause(c.p.ctx)
+}
+
+func (c Conn) SyncWrite(p []byte) (int, error) {
+	c.p.syncWriteLock.Lock()
+	defer c.p.syncWriteLock.Unlock()
+
+	c.p.writeCond.L.Lock()
+	// wait until buffer is exhausted
+	for c.p.writeStream.Len() != 0 && context.Cause(c.p.ctx) == nil {
+		c.p.writeCond.Wait()
+	}
+	c.p.writeStream.Write(p)
+	c.p.writeCond.L.Unlock()
+
+	if err := context.Cause(c.p.ctx); err != nil {
+		return len(p), err
+	}
+
+	c.p.writeCond.L.Lock()
+	// wait until data will be sent
+	for c.p.writeStream.Len() != 0 && context.Cause(c.p.ctx) == nil {
+		c.p.writeCond.Wait()
+	}
+	c.p.writeCond.L.Unlock()
 
 	return len(p), context.Cause(c.p.ctx)
 }
@@ -39,6 +69,8 @@ func (c Conn) Start() {
 }
 
 func (c Conn) start() {
+	defer c.p.writeCond.Broadcast()
+
 	buf := [tls.MaxRecordSize]byte{}
 
 	for {
@@ -48,9 +80,9 @@ func (c Conn) start() {
 		case <-c.p.clock.tick:
 		}
 
-		c.p.writeLock.Lock()
+		c.p.writeCond.L.Lock()
 		n, err := c.p.writeStream.Read(buf[:c.p.clock.stats.Size()])
-		c.p.writeLock.Unlock()
+		c.p.writeCond.L.Unlock()
 
 		if n == 0 || err != nil {
 			continue
@@ -60,6 +92,8 @@ func (c Conn) start() {
 			c.p.ctxCancel(err)
 			return
 		}
+
+		c.p.writeCond.Signal()
 	}
 }
 
@@ -75,6 +109,7 @@ func NewConn(ctx context.Context, conn essentials.Conn, stats *Stats) Conn {
 		p: &connPayload{
 			ctx:       ctx,
 			ctxCancel: cancel,
+			writeCond: sync.NewCond(&sync.Mutex{}),
 			clock: Clock{
 				stats: stats,
 				tick:  make(chan struct{}),