From 98d4cf7585de7381c8a85fd6666f5dca18c088dd Mon Sep 17 00:00:00 2001
From: xushiwei <x@goplus.org>
Date: Tue, 2 Jul 2024 20:23:48 +0800
Subject: [PATCH] runtime: close(chan)

---
 _demo/cchan/cchan.go       | 22 ++++++++++++++++++++++
 internal/runtime/z_chan.go | 38 +++++++++++++++++++++++++++++++++-----
 2 files changed, 55 insertions(+), 5 deletions(-)

diff --git a/_demo/cchan/cchan.go b/_demo/cchan/cchan.go
index f7c91b79..3d2e10a6 100644
--- a/_demo/cchan/cchan.go
+++ b/_demo/cchan/cchan.go
@@ -24,4 +24,26 @@ func doChan(cap int) {
 func main() {
 	doChan(10)
 	doChan(0)
+
+	c := runtime.NewChan(eltSize, 3)
+
+	v := 1
+	runtime.ChanSend(c, unsafe.Pointer(&v), eltSize)
+	v = 2
+	runtime.ChanSend(c, unsafe.Pointer(&v), eltSize)
+	v = 3
+	runtime.ChanSend(c, unsafe.Pointer(&v), eltSize)
+	runtime.ChanClose(c)
+
+	v = 10
+	if runtime.ChanTrySend(c, unsafe.Pointer(&v), eltSize) {
+		println("error: chan send to closed chan")
+	}
+
+	for {
+		if ok := runtime.ChanRecv(c, unsafe.Pointer(&v), eltSize); !ok {
+			break
+		}
+		println(v)
+	}
 }
diff --git a/internal/runtime/z_chan.go b/internal/runtime/z_chan.go
index e64b45ca..77fd957b 100644
--- a/internal/runtime/z_chan.go
+++ b/internal/runtime/z_chan.go
@@ -25,7 +25,9 @@ import (
 
 // -----------------------------------------------------------------------------
 
-const chanFull = 1
+const (
+	chanFull = 1
+)
 
 type Chan struct {
 	mutex sync.Mutex
@@ -34,6 +36,7 @@ type Chan struct {
 	getp  int
 	len   int
 	cap   int
+	close bool
 }
 
 func NewChan(eltSize, cap int) *Chan {
@@ -57,18 +60,25 @@ func ChanCap(p *Chan) int {
 	return p.cap
 }
 
+func ChanClose(p *Chan) {
+	p.mutex.Lock()
+	p.close = true
+	p.mutex.Unlock()
+	p.cond.Broadcast()
+}
+
 func ChanTrySend(p *Chan, v unsafe.Pointer, eltSize int) bool {
 	n := p.cap
 	p.mutex.Lock()
 	if n == 0 {
-		if p.getp == chanFull {
+		if p.getp == chanFull || p.close {
 			p.mutex.Unlock()
 			return false
 		}
 		p.data = v
 		p.getp = chanFull
 	} else {
-		if p.len == n {
+		if p.len == n || p.close {
 			p.mutex.Unlock()
 			return false
 		}
@@ -81,25 +91,34 @@ func ChanTrySend(p *Chan, v unsafe.Pointer, eltSize int) bool {
 	return true
 }
 
-func ChanSend(p *Chan, v unsafe.Pointer, eltSize int) {
+func ChanSend(p *Chan, v unsafe.Pointer, eltSize int) bool {
 	n := p.cap
 	p.mutex.Lock()
 	if n == 0 {
 		for p.getp == chanFull {
 			p.cond.Wait(&p.mutex)
 		}
+		if p.close {
+			p.mutex.Unlock()
+			return false
+		}
 		p.data = v
 		p.getp = chanFull
 	} else {
 		for p.len == n {
 			p.cond.Wait(&p.mutex)
 		}
+		if p.close {
+			p.mutex.Unlock()
+			return false
+		}
 		off := (p.getp + p.len) % n
 		c.Memcpy(c.Advance(p.data, off*eltSize), v, uintptr(eltSize))
 		p.len++
 	}
 	p.mutex.Unlock()
 	p.cond.Broadcast()
+	return true
 }
 
 func ChanTryRecv(p *Chan, v unsafe.Pointer, eltSize int) bool {
@@ -126,17 +145,25 @@ func ChanTryRecv(p *Chan, v unsafe.Pointer, eltSize int) bool {
 	return true
 }
 
-func ChanRecv(p *Chan, v unsafe.Pointer, eltSize int) {
+func ChanRecv(p *Chan, v unsafe.Pointer, eltSize int) bool {
 	n := p.cap
 	p.mutex.Lock()
 	if n == 0 {
 		for p.getp == 0 {
+			if p.close {
+				p.mutex.Unlock()
+				return false
+			}
 			p.cond.Wait(&p.mutex)
 		}
 		c.Memcpy(v, p.data, uintptr(eltSize))
 		p.getp = 0
 	} else {
 		for p.len == 0 {
+			if p.close {
+				p.mutex.Unlock()
+				return false
+			}
 			p.cond.Wait(&p.mutex)
 		}
 		c.Memcpy(v, c.Advance(p.data, p.getp*eltSize), uintptr(eltSize))
@@ -145,6 +172,7 @@ func ChanRecv(p *Chan, v unsafe.Pointer, eltSize int) {
 	}
 	p.mutex.Unlock()
 	p.cond.Broadcast()
+	return true
 }
 
 // -----------------------------------------------------------------------------