go-study/cyclic_barrier/h2o.go

package cyclic_barrier

import (
	"context"
	"fmt"
	"github.com/marusama/cyclicbarrier"
	"golang.org/x/sync/semaphore"
	"math/rand"
	"sync"
	"time"
)

const (
	H2OHydrogenNum = 2 // 氢原子数量
	H2OOxygenNum   = 1 // 氧原子数量
)

// H2O 水分子结构体
type H2O struct {
	semaH *semaphore.Weighted
	semaO *semaphore.Weighted
	cb    cyclicbarrier.CyclicBarrier
	wg    sync.WaitGroup
}

func NewH2O() *H2O {
	return &H2O{
		semaH: semaphore.NewWeighted(H2OHydrogenNum),            // 氢原子的信号量
		semaO: semaphore.NewWeighted(H2OOxygenNum),              // 氧原子的信号量
		cb:    cyclicbarrier.New(H2OHydrogenNum + H2OOxygenNum), // 循环栅栏，用来控制合成
	}
}

func (h *H2O) hydrogen(releaseHydrogen func()) {
	if err := h.semaH.Acquire(context.Background(), 1); err != nil { // 占用氢原子空槽
		fmt.Println("Hydrogen Acquire err:", err)
	}

	releaseHydrogen()
	if err := h.cb.Await(context.Background()); err != nil { // 等待栅栏放行
		fmt.Println("Hydrogen Await err:", err)
	}

	h.semaH.Release(1) // 释放氢原子空槽
}

func (h *H2O) oxygen(releaseOxygen func()) {
	if err := h.semaO.Acquire(context.Background(), 1); err != nil { // 占用氧原子空槽
		fmt.Println("Oxygen Acquire err:", err)
	}

	releaseOxygen()
	if err := h.cb.Await(context.Background()); err != nil { // 等待栅栏放行
		fmt.Println("Oxygen Await err:", err)
	}

	h.semaO.Release(1) // 释放氧原子空槽
}

func (h *H2O) Gen(num uint) <-chan string {
	// 计算总共需要生成的原子数量
	sum := H2OHydrogenNum + H2OOxygenNum
	numInt := int(num)
	// 用来存放生成的原子
	ch := make(chan string, numInt*sum)
	releaseHydrogen := func() {
		ch <- "H"
	}
	releaseOxygen := func() {
		ch <- "O"
	}

	// 使用 WaitGroup 等待所有的 goroutine 完成
	h.wg.Add(numInt * sum)
	for i := 0; i < numInt*H2OHydrogenNum; i++ {
		go func() {
			time.Sleep(time.Duration(rand.Intn(100)) * time.Millisecond)
			h.hydrogen(releaseHydrogen)
			h.wg.Done()
		}()
	}
	for i := 0; i < numInt*H2OOxygenNum; i++ {
		go func() {
			time.Sleep(time.Duration(rand.Intn(100)) * time.Millisecond)
			h.oxygen(releaseOxygen)
			h.wg.Done()
		}()
	}

	go func() {
		h.wg.Wait()
		close(ch)
	}()

	return ch
}