协程实现

May 06, 2020

from https://github.com/skywind3000/collection/tree/master/vintage/context

精简一部分代码 留下unix上使用的部分

首先是接口

#pragma once
#include <stddef.h>

// 切换需要保存到上下文
typedef struct _ctx_regs_t {
    unsigned long long rax;		// 0
    unsigned long long rbx;		// 8
    unsigned long long rcx;		// 16
    unsigned long long rdx;		// 24
    unsigned long long rsi;		// 32
    unsigned long long rdi;		// 40
    unsigned long long rsp;		// 48
    unsigned long long rbp;		// 56
    unsigned long long rip;		// 64
    unsigned long long rflags;  // 72
    unsigned long long r8;		// 80
    unsigned long long r9;		// 88
    unsigned long long r10;		// 96
    unsigned long long r11;		// 104
    unsigned long long r12;		// 112
    unsigned long long r13;		// 120
    unsigned long long r14;		// 128
    unsigned long long r15;		// 136
    unsigned int mxcsr;			// 144
    unsigned int padding;		// 148

    unsigned char fpregs[32];		// 44(32) / 152(64)
    unsigned int reserved[18];		// 76(32) / 186(64)
}	ctx_regs_t;

// 协程信息
typedef struct _ctx_context_t {
    struct _ctx_regs_t regs;
    struct _ctx_context_t *link;
    char *stack;
    int stack_size;
}	ctx_context_t;

// 保存当前上下文到ctx
int ctx_getcontext(ctx_context_t *ctx);

// 跳转至ctx的上下文
int ctx_setcontext(const ctx_context_t *ctx);

// 创建一个新协程 并保存在ctx内 fun为该协程的函数 args为参数
int ctx_makecontext(ctx_context_t *ctx, int (*fun)(void*args), void *args);

// 切换协程
// 旧的保存在save
// 把现在到上下文切去target
int ctx_swapcontext(ctx_context_t *save, const ctx_context_t *target);

下面是ctx_getcontext的实现,只需要保存上下文即可

.globl _ctx_getcontext
.globl ctx_getcontext

// 函数签名 int ctx_getcontext(ctx_context_t *ctx)
_ctx_getcontext:
ctx_getcontext:
	
    // 参数为rdi 为ctx 第一个参数
    movq %rdi, %rax 

    // 基本copy操作
    movq %rax, 0(%rax)
	movq %rbx, 8(%rax)
	movq %rcx, 16(%rax)
	movq %rdx, 24(%rax)
	movq %rsi, 32(%rax)
	movq %rdi, 40(%rax)


	leaq 8(%rsp), %rdx // 现在rdx是esp了
	movq %rdx, 48(%rax) // 保存ebp
	movq %rbp, 56(%rax) // 保存ebp 我们没有动ebp
    
    // 保存EIP
	movq 0(%rsp), %rdx  
	movq %rdx, 64(%rax) 

	pushfq  // push rflags 到栈
	popq %rdx   // 把rflags pop给rdx
	movq %rdx, 72(%rax) // 保存RDX
    
	// 以下为基本操作
	movq %r8, 80(%rax)
	movq %r9, 88(%rax)
	movq %r10, 96(%rax)
	movq %r11, 104(%rax)
	movq %r12, 112(%rax)
	movq %r13, 120(%rax)
	movq %r14, 128(%rax)
	movq %r15, 136(%rax)

	// 保存这两个
	// unsigned int mxcsr;
	// unsigned int padding;
	movq 24(%rax), %rdx
	stmxcsr 144(%rax)

	// unsigned char fpregs[32];
	// unsigned int reserved[18];
	fnstenv 152(%rax)
	fldenv 152(%rax)

	// 清空EAX
	xor %rax, %rax
	ret

对于ctx_setcontext 只需要把上下文搞回去

.globl _ctx_setcontext 
.globl ctx_setcontext 

// int ctx_getcontext(ctx_context_t *ctx);
_ctx_setcontext:   
ctx_setcontext:  

    // 参数为rdi 为ctx
	movq %rdi, %rax

    // 一个一个搞回去
	movq 8(%rax), %rbx
	movq 16(%rax), %rcx
	movq 24(%rax), %rdx
	movq 32(%rax), %rsi
	movq 40(%rax), %rdi
	movq 48(%rax), %rsp
	movq 56(%rax), %rbp
	movq 64(%rax), %rdx // 把rip搞对
    // 我没没法直接改rip地址
    // 只能通过ret来跳转 把rip push进栈 那么ret就会跳转这个地址
    // 假装这个是返回地址 也就是上一个ctx_getcontext的返回地址
	pushq %rdx 
	movq 72(%rax), %rdx
	pushq %rdx
	popfq
	movq 80(%rax), %r8
	movq 88(%rax), %r9
	movq 96(%rax), %r10
	movq 104(%rax), %r11
	movq 112(%rax), %r12
	movq 120(%rax), %r13
	movq 128(%rax), %r14
	movq 136(%rax), %r15
	movq 24(%rax), %rdx
	ldmxcsr 144(%rax)
	fldenv 152(%rax)
	// 返回1
	movq $1, %rax
	// 现在不是返回ctx_setcontext的调用者那里了
    // 而是跳去ctx那里
	ret

对于ctx_swapcontext 就是上面两个加起来

.globl _ctx_swapcontext
.globl ctx_swapcontext
_ctx_swapcontext:
ctx_swapcontext:
	movq %rdi, %rax 

	movq %rax, 0(%rax)
    movq %rbx, 8(%rax)
    movq %rcx, 16(%rax)
    movq %rdx, 24(%rax)
    movq %rsi, 32(%rax)
    movq %rdi, 40(%rax)

    
    leaq 8(%rsp), %rdx 
    movq %rdx, 48(%rax) 
    movq %rbp, 56(%rax) 
    movq 0(%rsp), %rdx  
    movq %rdx, 64(%rax) 

    pushfq  
    popq %rdx   
    movq %rdx, 72(%rax) 
    
    movq %r8, 80(%rax)
    movq %r9, 88(%rax)
    movq %r10, 96(%rax)
    movq %r11, 104(%rax)
    movq %r12, 112(%rax)
    movq %r13, 120(%rax)
    movq %r14, 128(%rax)
    movq %r15, 136(%rax)

    movq 24(%rax), %rdx
    stmxcsr 144(%rax)

    fnstenv 152(%rax)
    fldenv 152(%rax)

    // 保存完毕

	movq %rsi, %rax 

    movq 8(%rax), %rbx
    movq 16(%rax), %rcx
    movq 24(%rax), %rdx
    movq 32(%rax), %rsi
    movq 40(%rax), %rdi
    movq 48(%rax), %rsp
    movq 56(%rax), %rbp
    movq 64(%rax), %rdx 
    pushq %rdx 
    movq 72(%rax), %rdx
    pushq %rdx
    popfq
    movq 80(%rax), %r8
    movq 88(%rax), %r9
    movq 96(%rax), %r10
    movq 104(%rax), %r11
    movq 112(%rax), %r12
    movq 120(%rax), %r13
    movq 128(%rax), %r14
    movq 136(%rax), %r15
    movq 24(%rax), %rdx
    ldmxcsr 144(%rax)
    fldenv 152(%rax)
    movq $1, %rax
	ret

协程的构建

int ctx_makecontext(
ctx_context_t *ctx, 
int (*fun)(void*args), 
void *args)
{
	// 设置rsp到最后
	char *rsp = ctx->stack + ctx->stack_size;
	unsigned long long *stack;
	// 留位置放参数
    rsp = rsp - ((size_t)rsp & 63) - 64;
    // stack现在是最高的位置了
    stack = (unsigned long long*)rsp;
	stack[0] = 0;
	stack[1] = (unsigned long long)(ctx->link);
	stack[2] = (unsigned long long)((void*)fun);
	stack[3] = (unsigned long long)args;
    
    // 只要设置好rsp rip 使用setcontext即可随时起跳
	ctx->regs.rsp = (unsigned long long)rsp;
	ctx->regs.rip = (unsigned long long)((void*)ctx_context_startup);

	ctx->regs.rdi = stack[1]; // 1 parm
	ctx->regs.rsi = stack[2]; // 2 parm
	ctx->regs.rdx = stack[3]; // 3 parm
	return 0;
}

使用例子

#include "contextu.h"

ctx_context_t mc, cc;

int raw_thread(void*p) {
    printf("remote: hello %s\n", (char*)p);
    ctx_swapcontext(&cc, &mc);
    printf("remote: back again\n");
    ctx_swapcontext(&cc, &mc);
    printf("remote: return\n");
    return 0;
}

int main(void)
{
    char temp_stack[256];

    cc.stack = temp_stack;
    cc.stack_size = sizeof(temp_stack);
    cc.link = &mc;

    ctx_getcontext(&cc); // 把当前ctx保存在cc
    ctx_makecontext(&cc, raw_thread, (char*)"girl");

    printf("before switch: %d\n", cc.stack_size);

    ctx_swapcontext(&mc, &cc);

    printf("local: here\n");
    ctx_swapcontext(&mc, &cc);

    printf("local: again\n");
    ctx_swapcontext(&mc, &cc);

    printf("local: end\n");
    return 0;
}
Great! You've successfully subscribed.
Great! Next, complete checkout for full access.
Welcome back! You've successfully signed in.
Success! Your account is fully activated, you now have access to all content.