github.com/codingeasygo/util@v0.0.0-20231206062002-1ce2f004b7d9/xio/byte.go (about)

     1  package xio
     2  
     3  import (
     4  	"fmt"
     5  	"io"
     6  	"sync"
     7  )
     8  
     9  // ByteDistributeWriter is Writer by byte
    10  type ByteDistributeWriter struct {
    11  	ws  map[byte]io.Writer
    12  	lck sync.RWMutex
    13  }
    14  
    15  // NewByteDistributeWriter will return new ByteDistributeWriter
    16  func NewByteDistributeWriter() (writer *ByteDistributeWriter) {
    17  	writer = &ByteDistributeWriter{
    18  		ws:  map[byte]io.Writer{},
    19  		lck: sync.RWMutex{},
    20  	}
    21  	return
    22  }
    23  
    24  // Add WriteCloser to list
    25  func (h *ByteDistributeWriter) Add(m byte, w io.Writer) {
    26  	h.lck.Lock()
    27  	defer h.lck.Unlock()
    28  	h.ws[m] = w
    29  }
    30  
    31  func (h *ByteDistributeWriter) Write(b []byte) (n int, err error) {
    32  	h.lck.RLock()
    33  	writer, ok := h.ws[b[0]]
    34  	h.lck.RUnlock()
    35  	if ok {
    36  		n, err = writer.Write(b)
    37  	} else {
    38  		err = fmt.Errorf("writer not exist by %v", b[0])
    39  	}
    40  	return
    41  }
    42  
    43  // Close will close all connection
    44  func (h *ByteDistributeWriter) Close() (err error) {
    45  	h.lck.Lock()
    46  	for k, w := range h.ws {
    47  		if closer, ok := w.(io.Closer); ok {
    48  			closer.Close()
    49  		}
    50  		delete(h.ws, k)
    51  	}
    52  	h.lck.Unlock()
    53  	return
    54  }