From 0a603647b40e895b4e05829405717b538f927b2f Mon Sep 17 00:00:00 2001
From: Steven Allen <steven@stebalien.com>
Date: Wed, 22 Apr 2020 17:43:51 -0700
Subject: [PATCH] switch to storing sessions in the context

This patch allows associating a "session key" with a context by calling
`NewSession(ctx)`. When the session is used in a request, the exchange can call
`GetOrCreateSession(ctx)` to make the request in the appropriate session, or create a
new session with the given session key.

Implements https://github.com/ipfs/go-ipfs/issues/7198
---
 go.mod       |  2 ++
 interface.go | 70 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 72 insertions(+)

diff --git a/go.mod b/go.mod
index 4531116..c2667b9 100644
--- a/go.mod
+++ b/go.mod
@@ -4,3 +4,5 @@ require (
 	github.com/ipfs/go-block-format v0.0.2
 	github.com/ipfs/go-cid v0.0.5
 )
+
+go 1.13
diff --git a/interface.go b/interface.go
index c3032b2..8e24dbf 100644
--- a/interface.go
+++ b/interface.go
@@ -3,7 +3,9 @@ package exchange
 
 import (
 	"context"
+	"fmt"
 	"io"
+	"sync/atomic"
 
 	blocks "github.com/ipfs/go-block-format"
 	cid "github.com/ipfs/go-cid"
@@ -35,3 +37,71 @@ type SessionExchange interface {
 	Interface
 	NewSession(context.Context) Fetcher
 }
+
+// SessionID is an opaque type uniquely identifying a session.
+type SessionID struct {
+	// Opaque type to ensure users don't cook this up and break things.
+	id uint64
+}
+
+// String formats the session ID as a string.
+func (id SessionID) String() string {
+	return fmt.Sprintf("session-%d", id.id)
+}
+
+// IsZero returns true if this SessionId is the zero value (not a valid session
+// ID).
+func (id SessionID) IsZero() bool {
+	return id.id == 0
+}
+
+type sessionContextKey struct{}
+type sessionContextValue struct {
+	sesID  SessionID
+	sesCtx context.Context
+}
+
+// NewSession registers a new session with the context. The session will be
+// closed when the passed-in context is canceled.
+//
+// If there's already a session associated with the context, the existing
+// session will be used.
+//
+// This function does not initialize any state, it just reserves a new SessionID
+// associates it with the context.
+func NewSession(ctx context.Context) context.Context {
+	_, ctx = GetOrCreateSession(ctx)
+	return ctx
+}
+
+// last allocated session ID. 0 is _never_ used.
+var lastSessionID uint64
+
+// GetOrCreateSession loads the session from the context, or creates one if
+// there is no associated session.
+//
+// This function also returns the context used to create the session. The
+// session should be stopped when this context is canceled.
+func GetOrCreateSession(ctx context.Context) (SessionID, context.Context) {
+	if s, ok := ctx.Value(sessionContextKey{}).(*sessionContextValue); ok {
+		return s.sesID, s.sesCtx
+	}
+
+	// Allocate a new session ID
+	id := SessionID{atomic.AddUint64(&lastSessionID, 1)}
+
+	// Create a spot to spot to hold the session information.
+	ctxValue := &sessionContextValue{sesID: id}
+
+	// Derive a new context with this information.
+	sesCtx := context.WithValue(ctx, sessionContextKey{}, ctxValue)
+
+	// Cyclically reference the session context so the session's context
+	// also references the session.
+	//
+	// We could reference the original context, but that doesn't have the
+	// session attached to it.
+	ctxValue.sesCtx = sesCtx
+
+	return id, ctx
+}