From c7f10e1e67eff6ce2e3bd4d9279cd4a3d22b5430 Mon Sep 17 00:00:00 2001
From: Cholerae Hu <choleraehyq@gmail.com>
Date: Fri, 8 May 2020 11:15:12 +0800
Subject: [PATCH] pid: refactor and support arm64

Signed-off-by: Cholerae Hu <choleraehyq@gmail.com>
---
 pid_go1.5_amd64.go => pid_go1.5.go | 22 +---------
 pid_go1.5_amd64_test.go            | 16 -------
 pid_go1.5_arm64.s                  | 31 +++++++++++++
 pid_go1.5_test.go                  | 70 ++++++++++++++++++++++++++++++
 4 files changed, 102 insertions(+), 37 deletions(-)
 rename pid_go1.5_amd64.go => pid_go1.5.go (73%)
 delete mode 100644 pid_go1.5_amd64_test.go
 create mode 100644 pid_go1.5_arm64.s
 create mode 100644 pid_go1.5_test.go

diff --git a/pid_go1.5_amd64.go b/pid_go1.5.go
similarity index 73%
rename from pid_go1.5_amd64.go
rename to pid_go1.5.go
index a57e797..03beeb1 100644
--- a/pid_go1.5_amd64.go
+++ b/pid_go1.5.go
@@ -13,17 +13,11 @@
 // permissions and limitations under the License. See the AUTHORS file
 // for names of contributors.
 
-// +build amd64 amd64p32
+// +build amd64 amd64p32 arm64
 // +build gc,go1.5
 
 package goid
 
-import (
-	"unsafe"
-)
-
-var _ = unsafe.Sizeof(0)
-
 //go:nosplit
 func getPid() uintptr
 
@@ -31,17 +25,3 @@ func getPid() uintptr
 func GetPid() int {
 	return int(getPid())
 }
-
-//go:linkname procPin runtime.procPin
-//go:nosplit
-func procPin() int
-
-//go:linkname procUnpin runtime.procUnpin
-//go:nosplit
-func procUnpin()
-
-func getTID() int {
-	tid := procPin()
-	procUnpin()
-	return tid
-}
diff --git a/pid_go1.5_amd64_test.go b/pid_go1.5_amd64_test.go
deleted file mode 100644
index 5f7ced7..0000000
--- a/pid_go1.5_amd64_test.go
+++ /dev/null
@@ -1,16 +0,0 @@
-// +build amd64 amd64p32
-// +build gc,go1.5
-
-package goid
-
-import (
-	"testing"
-)
-
-func TestPid(t *testing.T) {
-	p1 := GetPid()
-	p2 := getTID()
-	if p1 != p2 {
-		t.Fatalf("The result of GetPid %d procPin %d are not equal!", p1, p2)
-	}
-}
diff --git a/pid_go1.5_arm64.s b/pid_go1.5_arm64.s
new file mode 100644
index 0000000..fa8e87d
--- /dev/null
+++ b/pid_go1.5_arm64.s
@@ -0,0 +1,31 @@
+// Copyright 2016 Peter Mattis.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+// implied. See the License for the specific language governing
+// permissions and limitations under the License. See the AUTHORS file
+// for names of contributors.
+
+// Assembly to mimic runtime.getg.
+
+// +build arm64
+// +build gc,go1.5
+
+#include "go_asm.h"
+#include "textflag.h"
+
+// func getPid() int64
+TEXT ·getPid(SB),NOSPLIT,$0-8
+    MOVD g, R14
+    MOVD g_m(R14), R13
+    MOVD m_p(R13), R14
+    MOVW p_id(R14), R13
+    MOVD R13, ret+0(FP)
+    RET
diff --git a/pid_go1.5_test.go b/pid_go1.5_test.go
new file mode 100644
index 0000000..747dbef
--- /dev/null
+++ b/pid_go1.5_test.go
@@ -0,0 +1,70 @@
+// +build gc,go1.5
+
+package goid
+
+import (
+	"fmt"
+	"testing"
+	"unsafe"
+)
+
+var _ = unsafe.Sizeof(0)
+
+//go:linkname procPin runtime.procPin
+//go:nosplit
+func procPin() int
+
+//go:linkname procUnpin runtime.procUnpin
+//go:nosplit
+func procUnpin()
+
+func getTID() int {
+	tid := procPin()
+	procUnpin()
+	return tid
+}
+
+func TestParallelGetPid(t *testing.T) {
+	ch := make(chan *string, 100)
+	for i := 0; i < cap(ch); i++ {
+		go func(i int) {
+			id := GetPid()
+			expected := getTID()
+			if id == expected {
+				ch <- nil
+				return
+			}
+			s := fmt.Sprintf("Expected %d, but got %d", expected, id)
+			ch <- &s
+		}(i)
+	}
+
+	for i := 0; i < cap(ch); i++ {
+		val := <-ch
+		if val != nil {
+			t.Fatal(*val)
+		}
+	}
+}
+
+func TestGetPid(t *testing.T) {
+	p1 := GetPid()
+	p2 := getTID()
+	if p1 != p2 {
+		t.Fatalf("The result of GetPid %d procPin %d are not equal!", p1, p2)
+	}
+}
+
+func BenchmarkGetPid(b *testing.B) {
+	for i := 0; i < b.N; i++ {
+		GetPid()
+	}
+}
+
+func BenchmarkParallelGetPid(b *testing.B) {
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			GetPid()
+		}
+	})
+}