From 7eee7c00a631f1debe36379b89974fb1fc5927be Mon Sep 17 00:00:00 2001
From: vyzo <vyzo@hackzen.org>
Date: Thu, 4 Feb 2021 00:03:12 +0200
Subject: [PATCH] add support for transient connections

---
 network/context.go | 17 +++++++++++++++++
 network/network.go |  2 ++
 2 files changed, 19 insertions(+)

diff --git a/network/context.go b/network/context.go
index 967feba..b9c95e1 100644
--- a/network/context.go
+++ b/network/context.go
@@ -13,9 +13,11 @@ var DialPeerTimeout = 60 * time.Second
 type noDialCtxKey struct{}
 type dialPeerTimeoutCtxKey struct{}
 type forceDirectDialCtxKey struct{}
+type useTransientCtxKey struct{}
 
 var noDial = noDialCtxKey{}
 var forceDirectDial = forceDirectDialCtxKey{}
+var useTransient = useTransientCtxKey{}
 
 // EXPERIMENTAL
 // WithForceDirectDial constructs a new context with an option that instructs the network
@@ -66,3 +68,18 @@ func GetDialPeerTimeout(ctx context.Context) time.Duration {
 func WithDialPeerTimeout(ctx context.Context, timeout time.Duration) context.Context {
 	return context.WithValue(ctx, dialPeerTimeoutCtxKey{}, timeout)
 }
+
+// WithUseTransient constructs a new context with an option that instructs to network
+// that it is acceptable to use a transient connection when opening a new stream.
+func WithUseTransient(ctx context.Context) context.Context {
+	return context.WithValue(ctx, useTransient, true)
+}
+
+// GetUseTransient returns true if the use transient option is set in the context.
+func GetUseTransient(ctx context.Context) bool {
+	v := ctx.Value(useTransient)
+	if v != nil {
+		return true
+	}
+	return false
+}
diff --git a/network/network.go b/network/network.go
index abaca40..fc1b30c 100644
--- a/network/network.go
+++ b/network/network.go
@@ -103,6 +103,8 @@ type Stat struct {
 	Direction Direction
 	// Opened is the timestamp when this connection was opened.
 	Opened time.Time
+	// Transient indicates that this connection is transient and may be closed soon
+	Transient bool
 	// Extra stores additional metadata about this connection.
 	Extra map[interface{}]interface{}
 }
-- 
GitLab