github.com/pdmccormick/importable-docker-buildx@v0.0.0-20240426161518-e47091289030/util/waitmap/waitmap.go (about)

     1  package waitmap
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  )
     7  
     8  type Map struct {
     9  	mu sync.RWMutex
    10  	m  map[string]interface{}
    11  	ch map[string]chan struct{}
    12  }
    13  
    14  func New() *Map {
    15  	return &Map{
    16  		m:  make(map[string]interface{}),
    17  		ch: make(map[string]chan struct{}),
    18  	}
    19  }
    20  
    21  func (m *Map) Set(key string, value interface{}) {
    22  	m.mu.Lock()
    23  	defer m.mu.Unlock()
    24  
    25  	m.m[key] = value
    26  
    27  	if ch, ok := m.ch[key]; ok {
    28  		if ch != nil {
    29  			close(ch)
    30  		}
    31  	}
    32  	m.ch[key] = nil
    33  }
    34  
    35  func (m *Map) Get(ctx context.Context, keys ...string) (map[string]interface{}, error) {
    36  	if len(keys) == 0 {
    37  		return map[string]interface{}{}, nil
    38  	}
    39  
    40  	if len(keys) > 1 {
    41  		out := make(map[string]interface{})
    42  		for _, key := range keys {
    43  			mm, err := m.Get(ctx, key)
    44  			if err != nil {
    45  				return nil, err
    46  			}
    47  			out[key] = mm[key]
    48  		}
    49  		return out, nil
    50  	}
    51  
    52  	key := keys[0]
    53  	m.mu.Lock()
    54  	ch, ok := m.ch[key]
    55  	if !ok {
    56  		ch = make(chan struct{})
    57  		m.ch[key] = ch
    58  	}
    59  
    60  	if ch != nil {
    61  		m.mu.Unlock()
    62  		select {
    63  		case <-ctx.Done():
    64  			return nil, ctx.Err()
    65  		case <-ch:
    66  			m.mu.Lock()
    67  		}
    68  	}
    69  
    70  	res := m.m[key]
    71  	m.mu.Unlock()
    72  
    73  	return map[string]interface{}{key: res}, nil
    74  }