github.com/blend/go-sdk@v1.20220411.3/bufferutil/buffer_handlers.go (about)

     1  /*
     2  
     3  Copyright (c) 2022 - Present. Blend Labs, Inc. All rights reserved
     4  Use of this source code is governed by a MIT license that can be found in the LICENSE file.
     5  
     6  */
     7  
     8  package bufferutil
     9  
    10  import (
    11  	"context"
    12  	"sync"
    13  
    14  	"github.com/blend/go-sdk/async"
    15  )
    16  
    17  // BufferHandlers is a synchronized map of listeners for new lines to a line buffer.
    18  type BufferHandlers struct {
    19  	sync.RWMutex
    20  	Handlers map[string]*async.Worker
    21  }
    22  
    23  // Add adds a handler.
    24  func (bl *BufferHandlers) Add(uid string, handler BufferChunkHandler) {
    25  	bl.Lock()
    26  	defer bl.Unlock()
    27  
    28  	if bl.Handlers == nil {
    29  		bl.Handlers = make(map[string]*async.Worker)
    30  	}
    31  
    32  	w := async.NewWorker(func(_ context.Context, wi interface{}) error {
    33  		handler(wi.(BufferChunk))
    34  		return nil
    35  	})
    36  	w.Work = make(chan interface{}, 128)
    37  
    38  	go func() { _ = w.Start() }()
    39  
    40  	bl.Handlers[uid] = w
    41  }
    42  
    43  // Remove removes a listener.
    44  func (bl *BufferHandlers) Remove(uid string) {
    45  	bl.Lock()
    46  	defer bl.Unlock()
    47  
    48  	if bl.Handlers == nil {
    49  		bl.Handlers = make(map[string]*async.Worker)
    50  	}
    51  	if w, ok := bl.Handlers[uid]; ok {
    52  		stopped := w.NotifyStopped()
    53  		_ = w.Stop()
    54  		<-stopped
    55  	}
    56  	delete(bl.Handlers, uid)
    57  }
    58  
    59  // Handle calls the handlers.
    60  func (bl *BufferHandlers) Handle(chunk BufferChunk) {
    61  	bl.RLock()
    62  	defer bl.RUnlock()
    63  
    64  	for _, handler := range bl.Handlers {
    65  		handler.Work <- chunk
    66  	}
    67  }
    68  
    69  // Close shuts down the handlers.
    70  func (bl *BufferHandlers) Close() {
    71  	bl.Lock()
    72  	defer bl.Unlock()
    73  
    74  	for _, queue := range bl.Handlers {
    75  		stopped := queue.NotifyStopped()
    76  		_ = queue.Stop()
    77  		<-stopped
    78  	}
    79  }